{% import "%namespace/%name_%version/%sub/services/%service/_shared_macros.j2" as shared_macros %}
{% extends '_base.py.j2' %}

{% block content %}

{% for method in service.methods.values() | selectattr('paged_result_field') %}
{% if loop.first %}
{# This lives within the loop in order to ensure that this template
   is empty if there are no paged methods.
 #}
from google.api_core import gapic_v1
from google.api_core import retry as retries
from google.api_core import retry_async as retries_async
from typing import Any, AsyncIterator, Awaitable, Callable, Sequence, Tuple, Optional, Iterator, Union
try:
    OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
    OptionalAsyncRetry = Union[retries_async.AsyncRetry, gapic_v1.method._MethodDefault, None]
except AttributeError:  # pragma: NO COVER
    OptionalRetry = Union[retries.Retry, object, None]  # type: ignore
    OptionalAsyncRetry = Union[retries_async.AsyncRetry, object, None]  # type: ignore

{% filter sort_lines %}
{% for method in service.methods.values() | selectattr('paged_result_field') %}
{{ method.input.ident.python_import }}
{{ method.output.ident.python_import }}
{% if not method.paged_result_field.is_primitive %}
{{ method.paged_result_field.message.ident.python_import }}
{% endif %}
{% endfor %}
{% endfilter %}
{% endif %}


class {{ method.name }}Pager:
    """A pager for iterating through ``{{ method.name|snake_case }}`` requests.

    This class thinly wraps an initial
    :class:`{{ method.output.ident.sphinx }}` object, and
    provides an ``__iter__`` method to iterate through its
    ``{{ method.paged_result_field.name }}`` field.

    If there are more pages, the ``__iter__`` method will make additional
    ``{{ method.name }}`` requests and continue to iterate
    through the ``{{ method.paged_result_field.name }}`` field on the
    corresponding responses.

    All the usual :class:`{{ method.output.ident.sphinx }}`
    attributes are available on the pager. If multiple requests are made, only
    the most recent response is retained, and thus used for attribute lookup.
    """
    def __init__(self,
            method: Callable[..., {{ method.output.ident }}],
            request: {{ method.input.ident }},
            response: {{ method.output.ident }},
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Union[float, object] = gapic_v1.method.DEFAULT,
            {{ shared_macros.client_method_metadata_argument()|indent(4) }} = {{ shared_macros.client_method_metadata_default_value() }}):
        """Instantiate the pager.

        Args:
            method (Callable): The method that was originally called, and
                which instantiated this pager.
            request ({{ method.input.ident.sphinx }}):
                The initial request object.
            response ({{ method.output.ident.sphinx }}):
                The initial response object.
            retry (google.api_core.retry.Retry): Designation of what errors,
                if any, should be retried.
            timeout (float): The timeout for this request.
            {{ shared_macros.client_method_metadata_argument_doc()|indent(4) }}
        """
        self._method = method
        self._request = {{ method.input.ident }}(request)
        self._response = response
        self._retry = retry
        self._timeout = timeout
        self._metadata = metadata

    def __getattr__(self, name: str) -> Any:
        return getattr(self._response, name)

    @property
    def pages(self) -> Iterator[{{ method.output.ident }}]:
        yield self._response
        while self._response.next_page_token:
            self._request.page_token = self._response.next_page_token
            self._response = self._method(self._request, retry=self._retry, timeout=self._timeout, metadata=self._metadata)
            yield self._response

    {% if method.paged_result_field.map %}
    def __iter__(self) -> Iterator[Tuple[str, {{ method.paged_result_field.type.fields.get('value').ident }}]]:
        for page in self.pages:
            yield from page.{{ method.paged_result_field.name}}.items()

    def get(self, key: str) -> Optional[{{ method.paged_result_field.type.fields.get('value').ident }}]:
        return self._response.{{ method.paged_result_field.name }}.get(key)
    {% else %}
    def __iter__(self) -> {{ method.paged_result_field.ident | replace('MutableSequence', 'Iterator') }}:
        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.

    This class thinly wraps an initial
    :class:`{{ method.output.ident.sphinx }}` object, and
    provides an ``__aiter__`` method to iterate through its
    ``{{ method.paged_result_field.name }}`` field.

    If there are more pages, the ``__aiter__`` method will make additional
    ``{{ method.name }}`` requests and continue to iterate
    through the ``{{ method.paged_result_field.name }}`` field on the
    corresponding responses.

    All the usual :class:`{{ method.output.ident.sphinx }}`
    attributes are available on the pager. If multiple requests are made, only
    the most recent response is retained, and thus used for attribute lookup.
    """
    def __init__(self,
            method: Callable[..., Awaitable[{{ method.output.ident }}]],
            request: {{ method.input.ident }},
            response: {{ method.output.ident }},
            *,
            retry: OptionalAsyncRetry = gapic_v1.method.DEFAULT,
            timeout: Union[float, object] = gapic_v1.method.DEFAULT,
            {{ shared_macros.client_method_metadata_argument()|indent(4) }} = {{ shared_macros.client_method_metadata_default_value() }}):
        """Instantiates the pager.

        Args:
            method (Callable): The method that was originally called, and
                which instantiated this pager.
            request ({{ method.input.ident.sphinx }}):
                The initial request object.
            response ({{ method.output.ident.sphinx }}):
                The initial response object.
            retry (google.api_core.retry.AsyncRetry): Designation of what errors,
                if any, should be retried.
            timeout (float): The timeout for this request.
            {{ shared_macros.client_method_metadata_argument_doc()|indent(4) }}
        """
        self._method = method
        self._request = {{ method.input.ident }}(request)
        self._response = response
        self._retry = retry
        self._timeout = timeout
        self._metadata = metadata

    def __getattr__(self, name: str) -> Any:
        return getattr(self._response, name)

    @property
    async def pages(self) -> AsyncIterator[{{ method.output.ident }}]:
        yield self._response
        while self._response.next_page_token:
            self._request.page_token = self._response.next_page_token
            self._response = await self._method(self._request, retry=self._retry, timeout=self._timeout, metadata=self._metadata)
            yield self._response
    {% if method.paged_result_field.map %}
    def __aiter__(self) -> Iterator[Tuple[str, {{ method.paged_result_field.type.fields.get('value').ident }}]]:
        async def async_generator():
            async for page in self.pages:
                for response in page.{{ method.paged_result_field.name}}.items():
                    yield response

        return async_generator()

    def get(self, key: str) -> Optional[{{ method.paged_result_field.type.fields.get('value').ident }}]:
        return self._response.{{ method.paged_result_field.name }}.get(key)

    {% else %}
    def __aiter__(self) -> {{ method.paged_result_field.ident | replace('MutableSequence', 'AsyncIterator') }}:
        async def async_generator():
            async for page in self.pages:
                for response in page.{{ method.paged_result_field.name }}:
                    yield response

        return async_generator()

    {% endif %}
    def __repr__(self) -> str:
        return '{0}<{1!r}>'.format(self.__class__.__name__, self._response)

{% endif %}
{% endfor %}
{% endblock %}
