Merge pull request googleapis#2602 from dhermes/pubsub-iterators

Using Iterators for list_topics() in Pub/Sub.

Author: dhermes

core/google/cloud/_testing.py

@@ -81,5 +81,12 @@ def __init__(self, items, page_token):
         self.page_token = page_token
 
     def next(self):
+        if self._items is None:
+            raise StopIteration
         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,16 +31,27 @@
 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):
     """Helper mapping publisher-related APIs.
 
     :type gax_api: :class:`google.pubsub.v1.publisher_api.PublisherApi`
     :param gax_api: API object used to make GAX requests.
+
+    :type client: :class:`~google.cloud.pubsub.client.Client`
+    :param client: The client that owns this API object.
     """
-    def __init__(self, gax_api):
+
+    def __init__(self, gax_api, client):
         self._gax_api = gax_api
+        self._client = client
 
     def list_topics(self, project, page_size=0, page_token=None):
         """List topics for the project associated with this API.
@@ -58,21 +71,21 @@ 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)
+
+        return Iterator(client=self._client, 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 +556,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

@@ -87,9 +87,9 @@ def publisher_api(self):
         if self._publisher_api is None:
             if self._use_gax:
                 generated = make_gax_publisher_api(self.connection)
-                self._publisher_api = GAXPublisherAPI(generated)
+                self._publisher_api = GAXPublisherAPI(generated, self)
             else:
-                self._publisher_api = JSONPublisherAPI(self.connection)
+                self._publisher_api = JSONPublisherAPI(self)
         return self._publisher_api
 
     @property
@@ -131,18 +131,13 @@ 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(
+        return api.list_topics(
             self.project, page_size, page_token)
-        topics = [Topic.from_api_repr(resource, self)
-                  for resource in resources]
-        return topics, next_token
 
     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'
@@ -97,12 +99,13 @@ def build_api_url(self, path, query_params=None,
 class _PublisherAPI(object):
     """Helper mapping publisher-related APIs.
 
-    :type connection: :class:`Connection`
-    :param connection: the connection used to make API requests.
+    :type client: :class:`~google.cloud.pubsub.client.Client`
+    :param client: the client used to make API requests.
     """
 
-    def __init__(self, connection):
-        self._connection = connection
+    def __init__(self, client):
+        self._client = client
+        self._connection = client.connection
 
     def list_topics(self, project, page_size=None, page_token=None):
         """API call:  list topics for a given project
@@ -122,24 +125,18 @@ 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')
+
+        return Iterator(client=self._client, 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 +573,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)