add support for custom pollers and pagers

autorest/codegen/models/lro_operation.py

@@ -83,15 +83,27 @@ def set_lro_response_type(self) -> None:
 
     @property
     def has_optional_return_type(self) -> bool:
-        """An LROOperation will never have an optional return type, we will always return LROPoller[return type]"""
+        """An LROOperation will never have an optional return type, we will always return a poller"""
         return False
 
+    def get_poller_path(self, async_mode: bool) -> str:
+        custom_poller_extension = "x-python-custom-poller-{}sync".format("a" if async_mode else "")
+        default_poller = "azure.core.polling.{}LROPoller".format("Async" if async_mode else "")
+        return self.yaml_data["extensions"].get(custom_poller_extension, default_poller)
+
+    def get_poller(self, async_mode: bool) -> str:
+        return self.get_poller_path(async_mode).split(".")[-1]
+
     def imports(self, code_model, async_mode: bool) -> FileImport:
         file_import = super().imports(code_model, async_mode)
         file_import.add_from_import("typing", "Union", ImportType.STDLIB, TypingSection.CONDITIONAL)
+
+        poller_import_path = ".".join(self.get_poller_path(async_mode).split(".")[:-1])
+        poller = self.get_poller(async_mode)
+
+        file_import.add_from_import(poller_import_path, poller, ImportType.AZURECORE)
         if async_mode:
             file_import.add_from_import("typing", "Optional", ImportType.STDLIB, TypingSection.CONDITIONAL)
-            file_import.add_from_import("azure.core.polling", "AsyncLROPoller", ImportType.AZURECORE)
             file_import.add_from_import("azure.core.polling", "AsyncNoPolling", ImportType.AZURECORE)
             file_import.add_from_import("azure.core.polling", "AsyncPollingMethod", ImportType.AZURECORE)
             if code_model.options['azure_arm']:
@@ -103,7 +115,6 @@ def imports(self, code_model, async_mode: bool) -> FileImport:
                     "azure.core.polling.async_base_polling", "AsyncLROBasePolling", ImportType.AZURECORE
                 )
         else:
-            file_import.add_from_import("azure.core.polling", "LROPoller", ImportType.AZURECORE)
             file_import.add_from_import("azure.core.polling", "NoPolling", ImportType.AZURECORE)
             file_import.add_from_import("azure.core.polling", "PollingMethod", ImportType.AZURECORE)
             if code_model.options['azure_arm']:

autorest/codegen/models/paging_operation.py

@@ -104,9 +104,19 @@ def next_link_name(self) -> Optional[str]:
 
     @property
     def has_optional_return_type(self) -> bool:
-        """A paging will never have an optional return type, we will always return ItemPaged[return type]"""
+        """A paging will never have an optional return type, we will always return a pager"""
         return False
 
+    def get_pager_path(self, async_mode: bool) -> str:
+        custom_pager_extension = "x-python-custom-pager-{}sync".format("a" if async_mode else "")
+        default_pager = "azure.core.{}paging.{}ItemPaged".format(
+            "async_" if async_mode else "", "Async" if async_mode else ""
+        )
+        return self.yaml_data["extensions"].get(custom_pager_extension, default_pager)
+
+    def get_pager(self, async_mode: bool) -> str:
+        return self.get_pager_path(async_mode).split(".")[-1]
+
     @property
     def success_status_code(self) -> List[Union[str, int]]:
         """The list of all successfull status code.
@@ -118,12 +128,15 @@ def success_status_code(self) -> List[Union[str, int]]:
     def imports(self, code_model, async_mode: bool) -> FileImport:
         file_import = super(PagingOperation, self).imports(code_model, async_mode)
 
+        pager_import_path = ".".join(self.get_pager_path(async_mode).split(".")[:-1])
+        pager = self.get_pager(async_mode)
+
+        file_import.add_from_import(pager_import_path, pager, ImportType.AZURECORE)
+
         if async_mode:
-            file_import.add_from_import("azure.core.async_paging", "AsyncItemPaged", ImportType.AZURECORE)
             file_import.add_from_import("azure.core.async_paging", "AsyncList", ImportType.AZURECORE)
             file_import.add_from_import("typing", "AsyncIterable", ImportType.STDLIB, TypingSection.CONDITIONAL)
         else:
-            file_import.add_from_import("azure.core.paging", "ItemPaged", ImportType.AZURECORE)
             file_import.add_from_import("typing", "Iterable", ImportType.STDLIB, TypingSection.CONDITIONAL)
 
         if code_model.options["tracing"]:

autorest/codegen/templates/lro_operation.py.jinja2

@@ -3,8 +3,8 @@
 {% set trace_decorator = "@distributed_trace_async" if async_mode else "@distributed_trace" %}
 {% set operation_name = "begin_"+operation.python_name %}
 {% macro return_docstring(async_mode) %}
-:return: An instance of {{ "Async" if async_mode }}LROPoller that returns either {{ operation.responses[0].docstring_text }} or the result of cls(response)
-:rtype: ~azure.core.polling.{{ "Async" if async_mode }}LROPoller[{{ operation.responses[0].docstring_type }}]{% endmacro %}
+:return: An instance of {{ operation.get_poller(async_mode) }} that returns either {{ operation.responses[0].docstring_text }} or the result of cls(response)
+:rtype: ~{{ operation.get_poller_path(async_mode)}}[{{ operation.responses[0].docstring_type }}]{% endmacro %}
 {% macro response_headers(response) %}
 response_headers = {
     {% for response_header in response.headers %}
@@ -21,7 +21,7 @@ response_headers = {
 {% if code_model.options['tracing'] %}
 {{ trace_decorator }}
 {% endif %}
-{% set return_type_wrapper = ["AsyncLROPoller" if async_mode else "LROPoller"] %}
+{% set return_type_wrapper = [operation.get_poller(async_mode)] %}
 {{ op_tools.method_signature(operation, operation_name, async_mode=async_mode, coroutine=async_mode, return_type_wrapper=return_type_wrapper) }}
 {%- if not async_mode %}
     {{ op_tools.sync_return_type_annotation(operation, return_type_wrapper) }}

autorest/codegen/templates/lro_operation_helper.jinja2

@@ -58,7 +58,6 @@
     {% set async_prefix = "Async" if async_mode else "" %}
     {% set path_format_arguments = "" %}
     {% set lro_options = (", lro_options={'final-state-via': '"+ operation.lro_options['final-state-via'] + "'}") if operation.lro_options else "" %}
-    {% set poller = "AsyncLROPoller" if async_mode else "LROPoller" %}
     {% set operation_name = "begin_"+operation.python_name %}
     {% if operation.parameters.path %}
     {% set path_format_arguments = ", path_format_arguments=path_format_arguments" %}
@@ -77,13 +76,13 @@
     elif polling is False: polling_method = {{ async_prefix }}NoPolling()
     else: polling_method = polling
     if cont_token:
-        return {{ poller }}.from_continuation_token(
+        return {{ operation.get_poller(async_mode) }}.from_continuation_token(
             polling_method=polling_method,
             continuation_token=cont_token,
             client=self._client,
             deserialization_callback=get_long_running_output
         )
     else:
-        return {{ poller }}(self._client, raw_result, get_long_running_output, polling_method)
+        return {{ operation.get_poller(async_mode) }}(self._client, raw_result, get_long_running_output, polling_method)
 {{ operation_name }}.metadata = {'url': '{{ operation.url }}'}  # type: ignore
 {%- endmacro -%}

autorest/codegen/templates/lro_paging_operation.py.jinja2

@@ -5,8 +5,8 @@
 {% set trace_decorator = "@distributed_trace_async" if async_mode else "@distributed_trace" %}
 {% set operation_name = "begin_"+operation.python_name %}
 {% macro return_docstring(async_mode) %}
-:return: An instance of {{ "Async" if async_mode }}LROPoller that returns an iterator like instance of either {{ operation.responses[0].docstring_text }} or the result of cls(response)
-:rtype: ~azure.core.polling.{{ "Async" if async_mode }}LROPoller[~azure.core.{{ "async_" if async_mode else "" }}paging.{{ "Async" if async_mode }}ItemPaged[{{ operation.responses[0].docstring_type }}]]{% endmacro %}
+:return: An instance of {{ operation.get_poller(async_mode) }} that returns an iterator like instance of either {{ operation.responses[0].docstring_text }} or the result of cls(response)
+:rtype: ~{{ operation.get_poller_path(async_mode) }}[~{{ operation.get_pager_path(async_mode) }}[{{ operation.responses[0].docstring_type }}]]{% endmacro %}
 {% macro operation_docstring(async_mode) %}
 {{ lro_helper.operation_docstring_helper(operation, async_mode) }}
 {{ return_docstring(async_mode) }}
@@ -16,7 +16,7 @@
 {% if code_model.options['tracing'] %}
 {{ trace_decorator }}
 {% endif %}
-{% set return_type_wrapper = [("Async" if async_mode else "") ~ "LROPoller", ("Async" if async_mode else "") ~ "ItemPaged"] %}
+{% set return_type_wrapper = [operation.get_poller(async_mode), operation.get_pager(async_mode)] %}
 {{ op_tools.method_signature(operation, operation_name, async_mode=async_mode, coroutine=async_mode, return_type_wrapper=return_type_wrapper) }}
 {%- if not async_mode %}
     {{ op_tools.sync_return_type_annotation(operation, return_type_wrapper) }}
@@ -32,7 +32,7 @@
             else:
                 return {{ keywords.await }}get_next(next_link)
 
-        return {{ "Async" if async_mode }}ItemPaged(
+        return {{ operation.get_pager(async_mode) }}(
             internal_get_next, extract_data
         )
 {{ lro_helper.lro_operation_return(code_model, operation, async_mode) }}

autorest/codegen/templates/metadata.json.jinja2

@@ -61,13 +61,13 @@
             "sync": {
                 {% if is_lro(operation) and is_paging(operation) %}
                     {% from "lro_paging_operation.py.jinja2" import operation_docstring with context %}
-                    {% set sync_return_type_wrapper = ["LROPoller", "ItemPaged"] %}
+                    {% set sync_return_type_wrapper = [operation.get_poller(async_mode=False), operation.get_pager(async_mode=False)] %}
                 {% elif is_lro(operation) %}
                     {% from "lro_operation.py.jinja2" import operation_docstring with context %}
-                    {% set sync_return_type_wrapper = ["LROPoller"] %}
+                    {% set sync_return_type_wrapper = [operation.get_poller(async_mode=False)] %}
                 {% elif is_paging(operation) %}
                     {% from "paging_operation.py.jinja2" import operation_docstring with context %}
-                    {% set sync_return_type_wrapper = ["ItemPaged"] %}
+                    {% set sync_return_type_wrapper = [operation.get_pager(async_mode=False)] %}
                 {% else %}
                     {% from "operation.py.jinja2" import operation_docstring with context %}
                     {% set sync_return_type_wrapper = "" %}
@@ -80,13 +80,13 @@
                 "coroutine": {{ coroutine | tojson }},
                 {% if is_lro(operation) and is_paging(operation) %}
                     {% from "lro_paging_operation.py.jinja2" import operation_docstring with context %}
-                    {% set async_return_type_wrapper = ["AsyncLROPoller", "AsyncItemPaged"] %}
+                    {% set async_return_type_wrapper = [operation.get_poller(async_mode=True), operation.get_pager(async_mode=True)] %}
                 {% elif is_lro(operation) %}
                     {% from "lro_operation.py.jinja2" import operation_docstring with context %}
-                    {% set async_return_type_wrapper = ["AsyncLROPoller"] %}
+                    {% set async_return_type_wrapper = [operation.get_poller(async_mode=True)] %}
                 {% elif is_paging(operation) %}
                     {% from "paging_operation.py.jinja2" import operation_docstring with context %}
-                    {% set async_return_type_wrapper = ["AsyncItemPaged"] %}
+                    {% set async_return_type_wrapper = [operation.get_pager(async_mode=True)] %}
                 {% else %}
                     {% from "operation.py.jinja2" import operation_docstring with context %}
                     {% set async_return_type_wrapper = "" %}

autorest/codegen/templates/paging_operation.py.jinja2

@@ -2,12 +2,10 @@
 {% import 'paging_operation_helper.jinja2' as helper %}
 {% set send_xml = "xml" if operation.parameters.has_body and "xml" in operation.request_content_type  %}
 {% set request_as_xml = ", is_xml=True" if send_xml else "" %}
-{% set item_paged = "AsyncItemPaged" if async_mode else "ItemPaged" %}
 {% macro return_docstring(async_mode) %}
 {% if operation.responses | selectattr('has_body') | first %}
 :return: An iterator like instance of either {{ operation.responses|selectattr('has_body')|map(attribute='docstring_text')|unique|join(' or ') }} or the result of cls(response)
-{# can't use item_paged variable, otherwise the call to get docstring from the metadata template will always return ItemPaged #}
-:rtype: ~azure.core.{{ "async_" if async_mode else "" }}paging.{{ "Async" if async_mode }}ItemPaged[{% for response in operation.responses %}{{response.docstring_type if response.has_body else "None"}}{% if not loop.last -%} or {% endif %}{% endfor %}]
+:rtype: ~{{ operation.get_pager_path(async_mode) }}[{% for response in operation.responses %}{{response.docstring_type if response.has_body else "None"}}{% if not loop.last -%} or {% endif %}{% endfor %}]
 {%- else -%}
 :return: None
 :rtype: None{%- endif -%}{%- endmacro -%}
@@ -54,7 +52,7 @@
 {% endif %}
     {{ helper.paging_operation(code_model, operation, async_mode) }}
 
-    return {{ item_paged }}(
+    return {{ operation.get_pager(async_mode) }}(
         get_next, extract_data
     )
 {{ operation.python_name }}.metadata = {'url': '{{ operation.url|replace("'", "\\'") }}'}  # type: ignore

tasks.py

@@ -311,6 +311,7 @@ def regenerate(c, swagger_name=None, debug=False):
         regenerate_multiapi(c, debug)
         regenerate_credential_default_policy(c, debug)
         regenerate_package_name_setup_py(c, debug)
+        regenerate_custom_poller_pager(c, debug)
 
 
 @task
@@ -386,3 +387,11 @@ def regenerate_multiapi(c, debug=False, swagger_name="test"):
         with Pool() as pool:
             result = pool.map(run_autorest, cmds)
         success = all(result)
+
+@task
+def regenerate_custom_poller_pager(c, debug=False):
+    cwd = os.getcwd()
+    cmd = (
+        f'{_AUTOREST_CMD_LINE} test/azure/specification/custompollerpager/README.md --use=. --python-sdks-folder={cwd}/test/'
+    )
+    success = run_autorest(cmd, debug=debug)

test/azure/Expected/AcceptanceTests/CustomPollerPager/custompollerpager/__init__.py

@@ -0,0 +1,19 @@
+# coding=utf-8
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# Code generated by Microsoft (R) AutoRest Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is regenerated.
+# --------------------------------------------------------------------------
+
+from ._auto_rest_paging_test_service import AutoRestPagingTestService
+from ._version import VERSION
+
+__version__ = VERSION
+__all__ = ['AutoRestPagingTestService']
+
+try:
+    from ._patch import patch_sdk  # type: ignore
+    patch_sdk()
+except ImportError:
+    pass

test/azure/Expected/AcceptanceTests/CustomPollerPager/custompollerpager/_auto_rest_paging_test_service.py

@@ -0,0 +1,67 @@
+# coding=utf-8
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# Code generated by Microsoft (R) AutoRest Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is regenerated.
+# --------------------------------------------------------------------------
+
+from typing import TYPE_CHECKING
+
+from azure.mgmt.core import ARMPipelineClient
+from msrest import Deserializer, Serializer
+
+if TYPE_CHECKING:
+    # pylint: disable=unused-import,ungrouped-imports
+    from typing import Any, Optional
+
+    from azure.core.credentials import TokenCredential
+
+from ._configuration import AutoRestPagingTestServiceConfiguration
+from .operations import PagingOperations
+from . import models
+
+
+class AutoRestPagingTestService(object):
+    """Long-running Operation for AutoRest.
+
+    :ivar paging: PagingOperations operations
+    :vartype paging: custompollerpager.operations.PagingOperations
+    :param credential: Credential needed for the client to connect to Azure.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :param str base_url: Service URL
+    :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
+    """
+
+    def __init__(
+        self,
+        credential,  # type: "TokenCredential"
+        base_url=None,  # type: Optional[str]
+        **kwargs  # type: Any
+    ):
+        # type: (...) -> None
+        if not base_url:
+            base_url = 'http://localhost:3000'
+        self._config = AutoRestPagingTestServiceConfiguration(credential, **kwargs)
+        self._client = ARMPipelineClient(base_url=base_url, config=self._config, **kwargs)
+
+        client_models = {k: v for k, v in models.__dict__.items() if isinstance(v, type)}
+        self._serialize = Serializer(client_models)
+        self._serialize.client_side_validation = False
+        self._deserialize = Deserializer(client_models)
+
+        self.paging = PagingOperations(
+            self._client, self._config, self._serialize, self._deserialize)
+
+    def close(self):
+        # type: () -> None
+        self._client.close()
+
+    def __enter__(self):
+        # type: () -> AutoRestPagingTestService
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *exc_details):
+        # type: (Any) -> None
+        self._client.__exit__(*exc_details)