Add wrap_with_paging

core/google/api/core/gapic_v1/method.py

@@ -18,10 +18,13 @@
 pagination, and long-running operations to gRPC methods.
 """
 
+import functools
 import platform
 
 import pkg_resources
+import six
 
+from google.api.core import page_iterator
 from google.api.core import timeout
 from google.api.core.helpers import grpc_helpers
 
@@ -238,4 +241,40 @@ def get_topic(name, timeout=None):
     if metadata is not None:
         metadata = _prepare_metadata(metadata)
 
-    return _GapicCallable(func, default_retry, default_timeout, metadata)
+    return six.wraps(func)(
+        _GapicCallable(func, default_retry, default_timeout, metadata))
+
+
+def wrap_with_paging(
+        func, items_field, request_token_field, response_token_field):
+    """Wrap an RPC method to return a page iterator.
+
+    Args:
+        func (Callable): The RPC method. This should already have been
+            wrapped with common functionality using :func:`wrap_method`.
+            request (protobuf.Message): The request message.
+        items_field (str): The field in the response message that has the
+            items for the page.
+        request_token_field (str): The field in the request message used to
+            specify the page token.
+        response_token_field (str): The field in the response message that has
+            the token for the next page.
+
+    Returns:
+        Callable: Returns a callable that when invoked will call the RPC
+            method and return a
+            :class:`google.api.core.page_iterator.Iterator`.
+    """
+    @six.wraps(func)
+    def paged_method(request, **kwargs):
+        """Wrapper that invokes a method and returns a page iterator."""
+        iterator = page_iterator.GRPCIterator(
+            client=None,
+            method=functools.partial(func, **kwargs),
+            request=request,
+            items_field=items_field,
+            request_token_field=request_token_field,
+            response_token_field=response_token_field)
+        return iterator
+
+    return paged_method

core/tests/unit/api_core/gapic/test_method.py

@@ -18,6 +18,7 @@
 from google.api.core import retry
 from google.api.core import timeout
 import google.api.core.gapic_v1.method
+import google.api.core.page_iterator
 
 
 def test_wrap_method_basic():
@@ -144,3 +145,38 @@ def test_wrap_method_with_overriding_timeout_as_a_number():
 
     assert result == 42
     method.assert_called_once_with(timeout=22, metadata=mock.ANY)
+
+
+def test_wrap_with_paging():
+    page_one = mock.Mock(
+        spec=['items', 'page_token', 'next_page_token'],
+        items=[1, 2],
+        next_page_token='icanhasnextpls')
+    page_two = mock.Mock(
+        spec=['items', 'page_token', 'next_page_token'],
+        items=[3, 4],
+        next_page_token=None)
+    method = mock.Mock(
+        spec=['__call__', '__name__'], side_effect=(page_one, page_two))
+    method.__name__ = 'mockmethod'
+
+    wrapped_method = google.api.core.gapic_v1.method.wrap_with_paging(
+        method, 'items', 'page_token', 'next_page_token')
+
+    request = mock.Mock(spec=['page_token'], page_token=None)
+    result = wrapped_method(request, extra='param')
+
+    # Should return an iterator and should not have actually called the
+    # method yet.
+    assert isinstance(result, google.api.core.page_iterator.Iterator)
+    method.assert_not_called()
+    assert request.page_token is None
+
+    # Draining the iterator should call the method until no more pages are
+    # returned.
+    results = list(result)
+
+    assert results == [1, 2, 3, 4]
+    assert method.call_count == 2
+    method.assert_called_with(request, extra='param')
+    assert request.page_token == 'icanhasnextpls'