Merge pull request #8 from Lumen5/topic-more-sane-search

Search: Respect total result field rather than running until HTTP 400
Lumen5 · Oct 1, 2018 · 7604971 · 7604971
2 parents 61fac34 + d4caa10
commit 7604971
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 19 deletions.
diff --git a/pygetty/__init__.py b/pygetty/__init__.py
@@ -2,4 +2,4 @@
 
 __doc__ = 'Wrapper for Getty v3 API'
 __author__ = 'Josh Klar <josh@lumen5.com>'
-__version__ = '1.1.5'
+__version__ = '1.1.6'
diff --git a/pygetty/search.py b/pygetty/search.py
@@ -95,22 +95,19 @@ def search(
 
     params['fields'] = ','.join(new_fields)
 
-    while returned < max_results:
-        try:
-            page = _fetch_page(
-                page=page_num,
-                page_size=page_size,
-                query_params=params,
-                asset_type=asset_type,
-                search_type=search_type,
-                auth_token_manager=auth_token_manager,
-            )
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code == 400 and 'page must be equal to' in e.response.text:
-                logger.warning('Got page must be equal to error')
-                return
+    total_results = None
+
+    while returned < max_results if total_results is None else min(max_results, total_results):
+        page = _fetch_page(
+            page=page_num,
+            page_size=page_size,
+            query_params=params,
+            asset_type=asset_type,
+            search_type=search_type,
+            auth_token_manager=auth_token_manager,
+        )
 
-            raise
+        total_results = page.get('result_count')
 
         for asset in page[asset_type]:
             yield asset_formatters[asset_type](asset)
@@ -119,9 +116,6 @@ def search(
             if returned >= max_results:
                 return
 
-        if len(page[asset_type]) < page_size:
-            return
-
         page_num += 1