add support for custom pollers and pagers

ChangeLog.md

@@ -1,5 +1,14 @@
 # Change Log
 
+### 2020-xx-xx - 5.5.0
+Autorest core version: 3.0.6318
+
+Modelerfour version: 4.15.421
+
+**New Features**
+
+- Can now take in custom pollers and pagers through directives. This will override the defaults (`LROPoller` and `ItemPaged`, respectively). See [this readme](https://github.com/Azure/autorest.python/tree/autorestv3/test/azure/specification/custompollerpager) for the directive to use to override.  #821
+
 ### 2020-11-11 - 5.4.3
 Autorest core version: 3.0.6320
 

autorest/codegen/models/lro_operation.py

@@ -85,15 +85,26 @@ 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:
+        extension_name = "poller-async" if async_mode else "poller-sync"
+        return self.yaml_data["extensions"][extension_name]
+
+    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']:
@@ -105,7 +116,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,16 @@ 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:
+        extension_name = "pager-async" if async_mode else "pager-sync"
+        return self.yaml_data["extensions"][extension_name]
+
+    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 +125,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

autorest/namer/name_converter.py

@@ -57,6 +57,26 @@ def _convert_operation_groups(operation_groups: List[Dict[str, Any]], code_model
                         NameConverter._convert_language_default_python_case(parameter, pad_string=PadType.Parameter)
                 for response in operation.get("responses", []):
                     NameConverter._convert_language_default_python_case(response)
+                if operation.get("extensions"):
+                    NameConverter._convert_extensions(operation)
+
+    @staticmethod
+    def _convert_extensions(operation: Dict[str, Any]) -> None:
+        operation_extensions = operation["extensions"]
+        if operation_extensions.get('x-ms-pageable'):
+            operation["extensions"]["pager-sync"] = operation_extensions.get(
+                "x-python-custom-pager-sync", "azure.core.paging.ItemPaged"
+            )
+            operation["extensions"]["pager-async"] = operation_extensions.get(
+                "x-python-custom-pager-async", "azure.core.async_paging.AsyncItemPaged"
+            )
+        if operation_extensions.get("x-ms-long-running-operation"):
+            operation["extensions"]["poller-sync"] = operation_extensions.get(
+                "x-python-custom-poller-sync", "azure.core.polling.LROPoller"
+            )
+            operation["extensions"]["poller-async"] = operation_extensions.get(
+                "x-python-custom-poller-async", "azure.core.polling.AsyncLROPoller"
+            )
 
     @staticmethod
     def _convert_schemas(schemas: Dict[str, Any]) -> None:

tasks.py

@@ -312,6 +312,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
@@ -387,3 +388,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/AcceptanceTests/asynctests/test_custom_poller_pager.py

@@ -0,0 +1,52 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+from async_generator import yield_, async_generator
+from custompollerpager.aio import AutoRestPagingTestService
+from custompollerpagerdefinitions.aio import AsyncCustomPager, AsyncCustomPoller
+
+import pytest
+
+@pytest.fixture
+@async_generator
+async def client(credential, authentication_policy):
+    async with AutoRestPagingTestService(credential, host="host:3000", authentication_policy=authentication_policy) as client:
+        await yield_(client)
+
+
+@pytest.fixture
+def custom_url_client(credential, authentication_policy):
+    with AutoRestParameterizedHostTestPagingClient(credential, host="host:3000", authentication_policy=authentication_policy) as client:
+        yield client
+
+class TestPaging(object):
+    def test_custom_pager(self, client):
+        pager = client.paging.get_single_pages()
+        assert isinstance(pager, AsyncCustomPager)
+
+    @pytest.mark.asyncio
+    async def test_custom_poller(self, client):
+        poller = await client.paging.begin_get_multiple_pages_lro()
+        assert isinstance(poller, AsyncCustomPoller)