Fixing crash on Python 3 when batch_size is set to None in the list p…

…ager (#203) Fixing crash on Python 3 when batch_size is set to None in the list pager
google · Mar 7, 2018 · da2e038 · da2e038
1 parent ac71613
commit da2e038
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 1 deletion.
diff --git a/apitools/base/py/list_pager.py b/apitools/base/py/list_pager.py
@@ -62,7 +62,16 @@ def YieldFromList(
     setattr(request, current_token_attribute, None)
     while limit is None or limit:
         if batch_size_attribute:
-            request_batch_size = min(batch_size, limit or batch_size)
+            # On Py3, None is not comparable so min() below will fail.
+            # On Py2, None is always less than any number so if batch_size
+            # is None, the request_batch_size will always be None regardless
+            # of the value of limit. This doesn't generally strike me as the
+            # correct behavior, but this change preserves the existing Py2
+            # behavior on Py3.
+            if batch_size is None:
+                request_batch_size = None
+            else:
+                request_batch_size = min(batch_size, limit or batch_size)
             setattr(request, batch_size_attribute, request_batch_size)
         response = getattr(service, method)(request,
                                             global_params=global_params)

diff --git a/apitools/base/py/list_pager_test.py b/apitools/base/py/list_pager_test.py
@@ -163,6 +163,36 @@ def testYieldFromListPaging(self):
 
         self._AssertInstanceSequence(results, 9)
 
+    def testYieldFromListBatchSizeNone(self):
+        self.mocked_client.column.List.Expect(
+            messages.FusiontablesColumnListRequest(
+                maxResults=None,
+                pageToken=None,
+                tableId='mytable',
+            ),
+            messages.ColumnList(
+                items=[
+                    messages.Column(name='c0'),
+                    messages.Column(name='c1'),
+                    messages.Column(name='c2'),
+                    messages.Column(name='c3'),
+                    messages.Column(name='c4'),
+                    messages.Column(name='c5'),
+                    messages.Column(name='c6'),
+                ],
+                nextPageToken='x',
+            ))
+
+        client = fusiontables.FusiontablesV1(get_credentials=False)
+        request = messages.FusiontablesColumnListRequest(tableId='mytable')
+        results = list_pager.YieldFromList(client.column,
+                                           request,
+                                           limit=5,
+                                           batch_size=None)
+
+        self._AssertInstanceSequence(results, 5)
+
+
     def testYieldFromListEmpty(self):
         self.mocked_client.column.List.Expect(
             messages.FusiontablesColumnListRequest(