fix: update paging implementation to handle unconventional pagination

gapic/schema/wrappers.py

@@ -866,13 +866,19 @@ def paged_result_field(self) -> Optional[Field]:
         """Return the response pagination field if the method is paginated."""
         # If the request field lacks any of the expected pagination fields,
         # then the method is not paginated.
-        for page_field in ((self.input, int, 'page_size'),
-                           (self.input, str, 'page_token'),
+        for page_field_token in ((self.input, str, 'page_token'),
                            (self.output, str, 'next_page_token')):
-            field = page_field[0].fields.get(page_field[2], None)
-            if not field or field.type != page_field[1]:
+            field = page_field_token[0].fields.get(page_field_token[2], None)
+            if not field or field.type != page_field_token[1]:
                 return None
 
+        page_fields = (self.input.fields.get('max_results', None),
+                       self.input.fields.get('page_size', None))
+        page_field_size = next(
+            (field for field in page_fields if field is not None), None)
+        if not page_field_size or page_field_size.type != int:
+            return None
+
         # Return the first repeated field.
         for field in self.output.fields.values():
             if field.repeated:

gapic/templates/%namespace/%name_%version/%sub/services/%service/pagers.py.j2

@@ -6,7 +6,7 @@
 {# This lives within the loop in order to ensure that this template
    is empty if there are no paged methods.
  -#}
-from typing import Any, AsyncIterable, Awaitable, Callable, Iterable, Sequence, Tuple
+from typing import Any, AsyncIterable, Awaitable, Callable, Iterable, Sequence, Tuple, Optional
 
 {% filter sort_lines -%}
 {% for method in service.methods.values() | selectattr('paged_result_field') -%}
@@ -68,14 +68,25 @@ class {{ method.name }}Pager:
             self._response = self._method(self._request, metadata=self._metadata)
             yield self._response
 
+    {% if method.paged_result_field.map %}
+    def __iter__(self) -> Iterable[Tuple[str, {{ method.paged_result_field.ident | replace('Sequence[', '') | replace(']', '') }}]]:
+        for page in self.pages:
+            yield from page.{{ method.paged_result_field.name}}.items()
+
+    def get(self, key: str) -> {{ method.paged_result_field.ident | replace('Sequence', 'Optional') }}:
+        return self._response.items.get(key)
+    {% else %}
     def __iter__(self) -> {{ method.paged_result_field.ident | replace('Sequence', 'Iterable') }}:
         for page in self.pages:
             yield from page.{{ method.paged_result_field.name }}
+    {% endif %}
 
     def __repr__(self) -> str:
         return '{0}<{1!r}>'.format(self.__class__.__name__, self._response)
 
 
+{# TODO(yon-mg): remove on rest async transport impl #}
+{% if 'grpc' in opts.transport %}
 class {{ method.name }}AsyncPager:
     """A pager for iterating through ``{{ method.name|snake_case }}`` requests.
 
@@ -138,5 +149,6 @@ class {{ method.name }}AsyncPager:
     def __repr__(self) -> str:
         return '{0}<{1!r}>'.format(self.__class__.__name__, self._response)
 
+{% endif %}
 {% endfor %}
 {% endblock %}

gapic/templates/noxfile.py.j2

@@ -32,6 +32,7 @@ def mypy(session):
     session.install('.')
     session.run(
         'mypy',
+        '--explicit-package-bases',
         {%- if api.naming.module_namespace %}
         '{{ api.naming.module_namespace[0] }}',
         {%- else %}

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -1020,7 +1020,7 @@ def test_{{ method.name|snake_case }}_raw_page_lro():
     assert response.raw_page is response
 {% endif %} {#- method.paged_result_field #}
 
-{% endfor -%} {#- method in methods #}
+{% endfor -%} {#- method in methods for grpc #}
 
 {% for method in service.methods.values() if 'rest' in opts.transport -%}
 def test_{{ method.name|snake_case }}_rest(transport: str = 'rest', request_type={{ method.input.ident }}):
@@ -1162,7 +1162,122 @@ def test_{{ method.name|snake_case }}_rest_flattened_error():
         )
 
 
-{% endfor -%}
+{% if method.paged_result_field %}
+def test_{{ method.name|snake_case }}_pager():
+    client = {{ service.client_name }}(
+        credentials=credentials.AnonymousCredentials(),
+    )
+
+    # Mock the http request call within the method and fake a response.
+    with mock.patch.object(Session, 'request') as req:
+        # Set the response as a series of pages
+        {% if method.paged_result_field.map%}
+        response = (
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}={
+                    'a':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                    'b':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                    'c':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                },
+                next_page_token='abc',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}={},
+                next_page_token='def',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}={
+                    'g':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                },
+                next_page_token='ghi',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}={
+                    'h':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                    'i':{{ method.paged_result_field.type.fields.get('value').ident }}(),
+                },
+            ),
+        )
+        {% else %}
+        response = (
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}=[
+                    {{ method.paged_result_field.type.ident }}(),
+                    {{ method.paged_result_field.type.ident }}(),
+                    {{ method.paged_result_field.type.ident }}(),
+                ],
+                next_page_token='abc',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}=[],
+                next_page_token='def',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}=[
+                    {{ method.paged_result_field.type.ident }}(),
+                ],
+                next_page_token='ghi',
+            ),
+            {{ method.output.ident }}(
+                {{ method.paged_result_field.name }}=[
+                    {{ method.paged_result_field.type.ident }}(),
+                    {{ method.paged_result_field.type.ident }}(),
+                ],
+            ),
+        )
+        {% endif %}
+        # Two responses for two calls
+        response = response + response
+
+        # Wrap the values into proper Response objs
+        response = tuple(map(lambda x: {{ method.output.ident }}.to_json(x), response))
+        side_effect = tuple(map(lambda x: Response(), response))
+        for return_val, response_val in zip(side_effect, response):
+            return_val._content = response_val.encode('UTF-8')
+        req.side_effect = side_effect
+
+        metadata = ()
+        {% if method.field_headers -%}
+        metadata = tuple(metadata) + (
+            gapic_v1.routing_header.to_grpc_metadata((
+                {%- for field_header in method.field_headers %}
+                {%- if not method.client_streaming %}
+                ('{{ field_header }}', ''),
+                {%- endif %}
+                {%- endfor %}
+            )),
+        )
+        {% endif -%}
+        pager = client.{{ method.name|snake_case }}(request={})
+
+        assert pager._metadata == metadata
+
+        {% if method.paged_result_field.map %}
+        assert isinstance(pager.get('a'), {{ method.paged_result_field.type.fields.get('value').ident }})
+        assert pager.get('h') is None
+        {% endif %}
+
+        results = [i for i in pager]
+        assert len(results) == 6
+        {% if method.paged_result_field.map %}
+        assert all(
+            isinstance(i, tuple) and
+            tuple(map(lambda x: type(x), results[0])) == (str, {{ method.paged_result_field.type.fields.get('value').ident }})
+                   for i in results)
+        assert pager.get('a') is None
+        assert isinstance(pager.get('h'), {{ method.paged_result_field.type.fields.get('value').ident }})
+        {% else %}
+        assert all(isinstance(i, {{ method.paged_result_field.type.ident }})
+                   for i in results)
+        {% endif %}
+
+        pages = list(client.{{ method.name|snake_case }}(request={}).pages)
+        for page_, token in zip(pages, ['abc','def','ghi', '']):
+            assert page_.raw_page.next_page_token == token
+
+
+{% endif %} {# paged methods #}
+{% endfor -%} {#- method in methods for rest #}
 def test_credentials_transport_error():
     # It is an error to provide credentials and a transport instance.
     transport = transports.{{ service.name }}{{ opts.transport[0].capitalize() }}Transport(

noxfile.py

@@ -227,7 +227,7 @@ def showcase_mypy(
         session.chdir(lib)
 
         # Run the tests.
-        session.run("mypy", "google")
+        session.run("mypy", "--explicit-package-bases", "google")
 
 
 @nox.session(python="3.8")

tests/unit/schema/wrappers/test_method.py

@@ -123,6 +123,16 @@ def test_method_paged_result_field_no_page_field():
                          )
     assert method.paged_result_field is None
 
+    method = make_method('Foo',
+                         input_message=make_message(name='FooRequest', fields=(
+                             make_field(name='page_token', type=9),  # str
+                         )),
+                         output_message=make_message(name='FooResponse', fields=(
+                             make_field(name='next_page_token', type=9),  # str
+                         ))
+                         )
+    assert method.paged_result_field is None
+
 
 def test_method_paged_result_ref_types():
     input_msg = make_message(
@@ -139,7 +149,7 @@ def test_method_paged_result_ref_types():
         name='ListMolluscsResponse',
         fields=(
             make_field(name='molluscs', message=mollusc_msg, repeated=True),
-            make_field(name='next_page_token', type=9)
+            make_field(name='next_page_token', type=9)  # str
         ),
         module='mollusc'
     )
@@ -207,7 +217,7 @@ def test_flattened_ref_types():
 
 
 def test_method_paged_result_primitive():
-    paged = make_field(name='squids', type=9, repeated=True)
+    paged = make_field(name='squids', type=9, repeated=True)    # str
     input_msg = make_message(
         name='ListSquidsRequest',
         fields=(