Using Iterators for list_topics() in Pub/Sub.

In the process, had to add custom support for the GAX
page iterator in our core Iterator implementation.

Author: dhermes

core/google/cloud/_testing.py

@@ -81,5 +81,14 @@ def __init__(self, items, page_token):
         self.page_token = page_token
 
     def next(self):
-        items, self._items = self._items, None
-        return items
+        if self._items is None:
+            raise StopIteration
+        else:
+            items = self._items
+            self._items = None
+            return items
+
+    __next__ = next
+
+    def __iter__(self):
+        return self

core/google/cloud/iterator.py

@@ -202,6 +202,14 @@ class Iterator(object):
                        takes the :class:`Iterator` that started the page,
                        the :class:`Page` that was started and the dictionary
                        containing the page response.
+
+    :type page_iter: callable
+    :param page_iter: (Optional) Callable to produce a pages iterator from the
+                      current iterator. Assumed signature takes the
+                      :class:`Iterator` that started the page. By default uses
+                      the HTTP pages iterator. Meant to provide a custom
+                      way to create pages (potentially with a custom
+                      transport such as gRPC).
     """
 
     _PAGE_TOKEN = 'pageToken'
@@ -211,7 +219,7 @@ class Iterator(object):
     def __init__(self, client, path, item_to_value,
                  items_key=DEFAULT_ITEMS_KEY,
                  page_token=None, max_results=None, extra_params=None,
-                 page_start=_do_nothing_page_start):
+                 page_start=_do_nothing_page_start, page_iter=None):
         self._started = False
         self.client = client
         self.path = path
@@ -220,8 +228,14 @@ def __init__(self, client, path, item_to_value,
         self.max_results = max_results
         self.extra_params = extra_params
         self._page_start = page_start
+        self._page_iter = None
+        # Verify inputs / provide defaults.
         if self.extra_params is None:
             self.extra_params = {}
+        if page_iter is None:
+            self._page_iter = self._default_page_iter()
+        else:
+            self._page_iter = page_iter(self)
         self._verify_params()
         # The attributes below will change over the life of the iterator.
         self.page_number = 0
@@ -239,7 +253,7 @@ def _verify_params(self):
             raise ValueError('Using a reserved parameter',
                              reserved_in_use)
 
-    def _pages_iter(self):
+    def _default_page_iter(self):
         """Generator of pages of API responses.
 
         Yields :class:`Page` instances.
@@ -263,11 +277,11 @@ def pages(self):
         if self._started:
             raise ValueError('Iterator has already started', self)
         self._started = True
-        return self._pages_iter()
+        return self._page_iter
 
     def _items_iter(self):
         """Iterator for each item returned."""
-        for page in self._pages_iter():
+        for page in self._page_iter:
             # Decrement the total results since the pages iterator adds
             # to it when each page is encountered.
             self.num_results -= page.num_items

core/unit_tests/test_iterator.py

@@ -128,6 +128,22 @@ def test_constructor_w_extra_param_collision(self):
         with self.assertRaises(ValueError):
             self._makeOne(client, path, None, extra_params=extra_params)
 
+    def test_constructor_non_default_page_iter(self):
+        connection = _Connection()
+        client = _Client(connection)
+        path = '/foo'
+        result = object()
+        called = []
+
+        def page_iter(iterator):
+            called.append(iterator)
+            return result
+
+        iterator = self._makeOne(client, path, None,
+                                 page_iter=page_iter)
+        self.assertEqual(called, [iterator])
+        self.assertIs(iterator._page_iter, result)
+
     def test_pages_iter_empty_then_another(self):
         import six
         from google.cloud._testing import _Monkey

docs/pubsub_snippets.py

@@ -46,13 +46,8 @@ def do_something_with(sub):  # pylint: disable=unused-argument
         pass
 
     # [START client_list_topics]
-    topics, token = client.list_topics()   # API request
-    while True:
-        for topic in topics:
-            do_something_with(topic)
-        if token is None:
-            break
-        topics, token = client.list_topics(page_token=token)  # API request
+    for topic in client.list_topics():   # API request(s)
+        do_something_with(topic)
     # [END client_list_topics]
 
 

pubsub/google/cloud/pubsub/_gax.py

@@ -14,6 +14,8 @@
 
 """GAX wrapper for Pubsub API requests."""
 
+import functools
+
 from google.cloud.gapic.pubsub.v1.publisher_api import PublisherApi
 from google.cloud.gapic.pubsub.v1.subscriber_api import SubscriberApi
 from google.gax import CallOptions
@@ -29,6 +31,12 @@
 from google.cloud._helpers import _pb_timestamp_to_rfc3339
 from google.cloud.exceptions import Conflict
 from google.cloud.exceptions import NotFound
+from google.cloud.iterator import Iterator
+from google.cloud.iterator import Page
+from google.cloud.pubsub.topic import Topic
+
+
+_FAKE_ITEMS_KEY = 'not-a-key'
 
 
 class _PublisherAPI(object):
@@ -58,21 +66,24 @@ def list_topics(self, project, page_size=0, page_token=None):
                            passed, the API will return the first page of
                            topics.
 
-        :rtype: tuple, (list, str)
-        :returns: list of ``Topic`` resource dicts, plus a
-                  "next page token" string:  if not None, indicates that
-                  more topics can be retrieved with another call (pass that
-                  value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of :class:`~google.cloud.pubsub.topic.Topic`
+                  accessible to the current API.
         """
         if page_token is None:
             page_token = INITIAL_PAGE
         options = CallOptions(page_token=page_token)
         path = 'projects/%s' % (project,)
         page_iter = self._gax_api.list_topics(
             path, page_size=page_size, options=options)
-        topics = [{'name': topic_pb.name} for topic_pb in page_iter.next()]
-        token = page_iter.page_token or None
-        return topics, token
+        page_iter = functools.partial(_recast_page_iterator, page_iter)
+
+        # NOTE: We don't currently have access to the client, so callers
+        #       that want the client, must manually bind the client to the
+        #       iterator instance returned.
+        return Iterator(client=None, path=path,
+                        item_to_value=_item_to_topic,
+                        page_iter=page_iter)
 
     def topic_create(self, topic_path):
         """API call:  create a topic
@@ -543,3 +554,42 @@ def make_gax_subscriber_api(connection):
     if connection.in_emulator:
         channel = insecure_channel(connection.host)
     return SubscriberApi(channel=channel)
+
+
+def _item_to_topic(iterator, resource):
+    """Convert a JSON job to the native object.
+
+    :type iterator: :class:`~google.cloud.iterator.Iterator`
+    :param iterator: The iterator that is currently in use.
+
+    :type resource: :class:`google.pubsub.v1.pubsub_pb2.Topic`
+    :param resource: A topic returned from the API.
+
+    :rtype: :class:`~google.cloud.pubsub.topic.Topic`
+    :returns: The next topic in the page.
+    """
+    return Topic.from_api_repr(
+        {'name': resource.name}, iterator.client)
+
+
+def _recast_page_iterator(page_iter, iterator):
+    """Wrap GAX pages generator.
+
+    In particular, wrap each page and capture some state from the
+    GAX iterator.
+
+    Yields :class:`~google.cloud.iterator.Page` instances
+
+    :type page_iter: :class:`~google.gax.PageIterator`
+    :param page_iter: The iterator to wrap.
+
+    :type iterator: :class:`~google.cloud.iterator.Iterator`
+    :param iterator: The iterator that owns each page.
+    """
+    for items in page_iter:
+        fake_response = {_FAKE_ITEMS_KEY: items}
+        page = Page(
+            iterator, fake_response, _FAKE_ITEMS_KEY, _item_to_topic)
+        iterator.next_page_token = page_iter.page_token or None
+        iterator.num_results += page.num_items
+        yield page

pubsub/google/cloud/pubsub/client.py

@@ -131,18 +131,17 @@ def list_topics(self, page_size=None, page_token=None):
                            passed, the API will return the first page of
                            topics.
 
-        :rtype: tuple, (list, str)
-        :returns: list of :class:`google.cloud.pubsub.topic.Topic`, plus a
-                  "next page token" string:  if not None, indicates that
-                  more topics can be retrieved with another call (pass that
-                  value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of :class:`~google.cloud.pubsub.topic.Topic`
+                  accessible to the current API.
         """
         api = self.publisher_api
-        resources, next_token = api.list_topics(
+        iterator = api.list_topics(
             self.project, page_size, page_token)
-        topics = [Topic.from_api_repr(resource, self)
-                  for resource in resources]
-        return topics, next_token
+        # NOTE: Make sure to set the client since ``api.list_topics()`` may
+        #       not have access to the current client.
+        iterator.client = self
+        return iterator
 
     def list_subscriptions(self, page_size=None, page_token=None):
         """List subscriptions for the project associated with this client.

pubsub/google/cloud/pubsub/connection.py

@@ -19,6 +19,8 @@
 
 from google.cloud import connection as base_connection
 from google.cloud.environment_vars import PUBSUB_EMULATOR
+from google.cloud.iterator import Iterator
+from google.cloud.pubsub.topic import Topic
 
 
 PUBSUB_API_HOST = 'pubsub.googleapis.com'
@@ -104,7 +106,8 @@ class _PublisherAPI(object):
     def __init__(self, connection):
         self._connection = connection
 
-    def list_topics(self, project, page_size=None, page_token=None):
+    @staticmethod
+    def list_topics(project, page_size=None, page_token=None):
         """API call:  list topics for a given project
 
         See:
@@ -122,24 +125,21 @@ def list_topics(self, project, page_size=None, page_token=None):
                            passed, the API will return the first page of
                            topics.
 
-        :rtype: tuple, (list, str)
-        :returns: list of ``Topic`` resource dicts, plus a
-                  "next page token" string:  if not None, indicates that
-                  more topics can be retrieved with another call (pass that
-                  value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of :class:`~google.cloud.pubsub.topic.Topic`
+                  accessible to the current connection.
         """
-        conn = self._connection
-        params = {}
-
+        extra_params = {}
         if page_size is not None:
-            params['pageSize'] = page_size
-
-        if page_token is not None:
-            params['pageToken'] = page_token
-
+            extra_params['pageSize'] = page_size
         path = '/projects/%s/topics' % (project,)
-        resp = conn.api_request(method='GET', path=path, query_params=params)
-        return resp.get('topics', ()), resp.get('nextPageToken')
+
+        # NOTE: We don't currently have access to the client, so callers
+        #       that want the client, must manually bind the client to the
+        #       iterator instance returned.
+        return Iterator(client=None, path=path,
+                        items_key='topics', item_to_value=_item_to_topic,
+                        page_token=page_token, extra_params=extra_params)
 
     def topic_create(self, topic_path):
         """API call:  create a topic
@@ -576,3 +576,18 @@ def _transform_messages_base64(messages, transform, key=None):
             message = message[key]
         if 'data' in message:
             message['data'] = transform(message['data'])
+
+
+def _item_to_topic(iterator, resource):
+    """Convert a JSON job to the native object.
+
+    :type iterator: :class:`~google.cloud.iterator.Iterator`
+    :param iterator: The iterator that is currently in use.
+
+    :type resource: dict
+    :param resource: A topic returned from the API.
+
+    :rtype: :class:`~google.cloud.pubsub.topic.Topic`
+    :returns: The next topic in the page.
+    """
+    return Topic.from_api_repr(resource, iterator.client)

pubsub/unit_tests/test__gax.py

@@ -56,17 +56,24 @@ def test_ctor(self):
     def test_list_topics_no_paging(self):
         from google.gax import INITIAL_PAGE
         from google.cloud._testing import _GAXPageIterator
+        from google.cloud.pubsub.topic import Topic
+
         TOKEN = 'TOKEN'
         response = _GAXPageIterator([_TopicPB(self.TOPIC_PATH)], TOKEN)
         gax_api = _GAXPublisherAPI(_list_topics_response=response)
         api = self._makeOne(gax_api)
 
-        topics, next_token = api.list_topics(self.PROJECT)
+        iterator = api.list_topics(self.PROJECT)
+        # Add back the client to support API requests.
+        iterator.client = _Client(self.PROJECT)
+        topics = list(iterator)
+        next_token = iterator.next_page_token
 
         self.assertEqual(len(topics), 1)
         topic = topics[0]
-        self.assertIsInstance(topic, dict)
-        self.assertEqual(topic['name'], self.TOPIC_PATH)
+        self.assertIsInstance(topic, Topic)
+        self.assertEqual(topic.name, self.TOPIC_NAME)
+        self.assertEqual(topic.full_name, self.TOPIC_PATH)
         self.assertEqual(next_token, TOKEN)
 
         name, page_size, options = gax_api._list_topics_called_with
@@ -76,6 +83,8 @@ def test_list_topics_no_paging(self):
 
     def test_list_topics_with_paging(self):
         from google.cloud._testing import _GAXPageIterator
+        from google.cloud.pubsub.topic import Topic
+
         SIZE = 23
         TOKEN = 'TOKEN'
         NEW_TOKEN = 'NEW_TOKEN'
@@ -84,13 +93,18 @@ def test_list_topics_with_paging(self):
         gax_api = _GAXPublisherAPI(_list_topics_response=response)
         api = self._makeOne(gax_api)
 
-        topics, next_token = api.list_topics(
+        iterator = api.list_topics(
             self.PROJECT, page_size=SIZE, page_token=TOKEN)
+        # Add back the client to support API requests.
+        iterator.client = _Client(self.PROJECT)
+        topics = list(iterator)
+        next_token = iterator.next_page_token
 
         self.assertEqual(len(topics), 1)
         topic = topics[0]
-        self.assertIsInstance(topic, dict)
-        self.assertEqual(topic['name'], self.TOPIC_PATH)
+        self.assertIsInstance(topic, Topic)
+        self.assertEqual(topic.name, self.TOPIC_NAME)
+        self.assertEqual(topic.full_name, self.TOPIC_PATH)
         self.assertEqual(next_token, NEW_TOKEN)
 
         name, page_size, options = gax_api._list_topics_called_with
@@ -1043,3 +1057,9 @@ class _Connection(object):
     def __init__(self, in_emulator=False, host=None):
         self.in_emulator = in_emulator
         self.host = host
+
+
+class _Client(object):
+
+    def __init__(self, project):
+        self.project = project