Using Iterators for list_topics() in Pub/Sub.

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)