Issue 119: Change retrieve_data to be a generator

8h2a · web-flow · commit e74765ba7f9f · 2019-07-03T23:01:00.000+02:00
See issue josegonzalez#119.
diff --git a/bin/github-backup b/bin/github-backup
@@ -406,12 +406,11 @@ def get_github_repo_url(args, repository):
     return repo_url
 
 
-def retrieve_data(args, template, query_args=None, single_request=False):
+def retrieve_data_gen(args, template, query_args=None, single_request=False):
     auth = get_auth(args)
     query_args = get_query_args(query_args)
     per_page = 100
     page = 0
-    data = []
 
     while True:
         page = page + 1
@@ -438,20 +437,21 @@ def retrieve_data(args, template, query_args=None, single_request=False):
         response = json.loads(r.read().decode('utf-8'))
         if len(errors) == 0:
             if type(response) == list:
-                data.extend(response)
+                for resp in response:
+                    yield resp
                 if len(response) < per_page:
                     break
             elif type(response) == dict and single_request:
-                data.append(response)
+                yield response
 
         if len(errors) > 0:
             log_error(errors)
 
         if single_request:
             break
 
-    return data
-
+def retrieve_data(args, template, query_args=None, single_request=False):
+    return list(retrieve_data_gen(args, template, query_args, single_request))
 
 def get_query_args(query_args=None):
     if not query_args:
@@ -836,18 +836,21 @@ def backup_pulls(args, repo_cwd, repository, repos_template):
         pull_states = ['open', 'closed']
         for pull_state in pull_states:
             query_args['state'] = pull_state
-            # It'd be nice to be able to apply the args.since filter here...
-            _pulls = retrieve_data(args,
+            _pulls = retrieve_data_gen(args,
                                    _pulls_template,
                                    query_args=query_args)
             for pull in _pulls:
+                if args.since and pull['updated_at'] < args.since:
+                    break
                 if not args.since or pull['updated_at'] >= args.since:
                     pulls[pull['number']] = pull
     else:
-        _pulls = retrieve_data(args,
+        _pulls = retrieve_data_gen(args,
                                _pulls_template,
                                query_args=query_args)
         for pull in _pulls:
+            if args.since and pull['updated_at'] < args.since:
+                break
             if not args.since or pull['updated_at'] >= args.since:
                 pulls[pull['number']] = retrieve_data(
                     args,
