fix default polling method listed in docstring

autorest/codegen/models/lro_operation.py

@@ -88,38 +88,66 @@ def has_optional_return_type(self) -> bool:
         """An LROOperation will never have an optional return type, we will always return a poller"""
         return False
 
+    def _get_lro_extension(self, extension_base, async_mode, *, azure_arm=None):
+        extension_name = extension_base + ("-async" if async_mode else "-sync")
+        extension = self.yaml_data["extensions"][extension_name]
+        arm_extension = None
+        if azure_arm is not None:
+            arm_extension = "azure-arm" if azure_arm else "data-plane"
+        return extension[arm_extension] if arm_extension else extension
+
     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]
+        return self._get_lro_extension("poller", async_mode)
 
     def get_poller(self, async_mode: bool) -> str:
         return self.get_poller_path(async_mode).split(".")[-1]
 
+    def get_default_polling_method_path(self, async_mode: bool, azure_arm: bool) -> str:
+        return self._get_lro_extension("default-polling-method", async_mode, azure_arm=azure_arm)
+
+    def get_default_polling_method(self, async_mode: bool, azure_arm: bool) -> str:
+        return self.get_default_polling_method_path(async_mode, azure_arm).split(".")[-1]
+
+    def get_default_no_polling_method_path(self, async_mode: bool) -> str:
+        return self._get_lro_extension("default-no-polling-method", async_mode)
+
+    def get_default_no_polling_method(self, async_mode: bool) -> str:
+        return self.get_default_no_polling_method_path(async_mode).split(".")[-1]
+
+    def get_base_polling_method_path(self, async_mode: bool) -> str:
+        return self._get_lro_extension("base-polling-method", async_mode)
+
+    def get_base_polling_method(self, async_mode: bool) -> str:
+        return self.get_base_polling_method_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)
+
+        default_polling_method_import_path = ".".join(
+            self.get_default_polling_method_path(async_mode, code_model.options['azure_arm']).split(".")[:-1]
+        )
+        default_polling_method = self.get_default_polling_method(async_mode, code_model.options['azure_arm'])
+        file_import.add_from_import(default_polling_method_import_path, default_polling_method, ImportType.AZURECORE)
+
+        default_no_polling_method_import_path = ".".join(
+            self.get_default_no_polling_method_path(async_mode).split(".")[:-1]
+        )
+        default_no_polling_method = self.get_default_no_polling_method(async_mode)
+        file_import.add_from_import(
+            default_no_polling_method_import_path, default_no_polling_method, ImportType.AZURECORE
+        )
+
+        base_polling_method_import_path = ".".join(
+            self.get_base_polling_method_path(async_mode).split(".")[:-1]
+        )
+        base_polling_method = self.get_base_polling_method(async_mode)
+        file_import.add_from_import(base_polling_method_import_path, base_polling_method, ImportType.AZURECORE)
+
         if async_mode:
             file_import.add_from_import("typing", "Optional", ImportType.STDLIB, TypingSection.CONDITIONAL)
-            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']:
-                file_import.add_from_import(
-                    "azure.mgmt.core.polling.async_arm_polling", "AsyncARMPolling", ImportType.AZURECORE
-                )
-            else:
-                file_import.add_from_import(
-                    "azure.core.polling.async_base_polling", "AsyncLROBasePolling", ImportType.AZURECORE
-                )
-        else:
-            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']:
-                file_import.add_from_import("azure.mgmt.core.polling.arm_polling", "ARMPolling", ImportType.AZURECORE)
-            else:
-                file_import.add_from_import("azure.core.polling.base_polling", "LROBasePolling", ImportType.AZURECORE)
         return file_import

autorest/codegen/models/paging_operation.py

@@ -80,6 +80,9 @@ def _find_python_name(self, rest_api_name: str, log_name: str) -> str:
             + f"{log_name}:{rest_api_name} in model {response_schema.name}"
         )
 
+    def _get_paging_extension(self, extension_name):
+        return self.yaml_data["extensions"][extension_name]
+
     @property
     def item_name(self) -> str:
         if self._item_name is None:
@@ -109,7 +112,7 @@ def has_optional_return_type(self) -> bool:
 
     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]
+        return self._get_paging_extension(extension_name)
 
     def get_pager(self, async_mode: bool) -> str:
         return self.get_pager_path(async_mode).split(".")[-1]

autorest/codegen/templates/keywords.jinja2

@@ -1,4 +1,5 @@
 {% set def = "async def" if async_mode else "def" %}
 {% set async_prefix = "a" if async_mode else "" %}
 {% set await = "await " if async_mode else "" %}
+{% set async_class = "Async" if async_mode else "" %}
 {% macro escape_str(s) %}'{{ s|replace("'", "\\'") }}'{% endmacro %}

autorest/codegen/templates/lro_operation.py.jinja2

@@ -13,7 +13,7 @@ response_headers = {
 }
 {% endmacro %}
 {% macro operation_docstring(async_mode) %}
-{{ helper.operation_docstring_helper(operation, async_mode) }}
+{{ helper.operation_docstring_helper(code_model, operation, async_mode) }}
 {{ return_docstring(async_mode) }}
 :raises ~azure.core.exceptions.HttpResponseError:
 """{% endmacro %}

autorest/codegen/templates/lro_operation_helper.jinja2

@@ -2,7 +2,8 @@
 
 {% macro param_documentation_string(parameter) %}:param {{ parameter.serialized_name }}: {{ parameter.description }}{% endmacro %}
 
-{% macro operation_docstring_helper(operation, async_mode) %}
+{% macro operation_docstring_helper(code_model, operation, async_mode) %}
+{% import 'keywords.jinja2' as keywords with context %}
 """{{ operation.summary if operation.summary else operation.description | wordwrap(width=95, break_long_words=False, wrapstring='\n') }}
 {% if operation.summary and operation.description %}
 
@@ -25,16 +26,16 @@
 {% endif %}
 :keyword callable cls: A custom type or function that will be passed the direct response
 :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-:keyword polling: True for ARMPolling, False for no polling, or a
- polling object for personal polling strategy
-:paramtype polling: bool or ~azure.core.polling.{{ "Async" if async_mode else "" }}PollingMethod
+:keyword polling: Pass in True if you'd like the {{ operation.get_default_polling_method(async_mode, code_model.options["azure_arm"]) }} polling method,
+ False for no polling, or your own initialized polling object for a personal polling strategy.
+:paramtype polling: bool or ~azure.core.polling.{{ keywords.async_class }}PollingMethod
 :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
 {%- endmacro -%}
 
 
 {% macro lro_operation(code_model, operation, async_mode) %}
 {% import 'keywords.jinja2' as keywords with context %}
-    polling = kwargs.pop('polling', {{ "True" if code_model.options['azure_arm'] else "False" }})  # type: Union[bool, {{ "Async" if async_mode else "" }}PollingMethod]
+    polling = kwargs.pop('polling', {{ "True" if code_model.options['azure_arm'] else "False" }})  # type: Union[bool, {{ keywords.async_class }}PollingMethod]
     cls = kwargs.pop('cls', None)  # type: ClsType[{{ op_tools.return_type_annotation(operation) }}]
     lro_delay = kwargs.pop(
         'polling_interval',
@@ -55,7 +56,7 @@
     {%- endmacro -%}
 
     {% macro lro_operation_return(code_model, operation, async_mode) %}
-    {% set async_prefix = "Async" if async_mode else "" %}
+    {% import 'keywords.jinja2' as keywords with context %}
     {% set path_format_arguments = "" %}
     {% set lro_options = (", lro_options={'final-state-via': '"+ operation.lro_options['final-state-via'] + "'}") if operation.lro_options else "" %}
     {% set operation_name = "begin_"+operation.python_name %}
@@ -68,12 +69,8 @@
     }
 
     {% endif %}
-    {% if code_model.options['azure_arm'] %}
-    if polling is True: polling_method = {{ async_prefix }}ARMPolling(lro_delay{{ lro_options }}{{ path_format_arguments }},  **kwargs)
-    {% else %}
-    if polling is True: polling_method = {{ async_prefix }}LROBasePolling(lro_delay{{ lro_options }}{{ path_format_arguments }},  **kwargs)
-    {% endif %}
-    elif polling is False: polling_method = {{ async_prefix }}NoPolling()
+    if polling is True: polling_method = {{ operation.get_default_polling_method(async_mode, code_model.options["azure_arm"]) }}(lro_delay{{ lro_options }}{{ path_format_arguments }},  **kwargs)
+    elif polling is False: polling_method = {{ operation.get_default_no_polling_method(async_mode) }}()
     else: polling_method = polling
     if cont_token:
         return {{ operation.get_poller(async_mode) }}.from_continuation_token(

autorest/codegen/templates/lro_paging_operation.py.jinja2

@@ -8,7 +8,7 @@
 :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) }}
+{{ lro_helper.operation_docstring_helper(code_model, operation, async_mode) }}
 {{ return_docstring(async_mode) }}
 :raises ~azure.core.exceptions.HttpResponseError:
 """{% endmacro %}

autorest/namer/name_converter.py

@@ -71,13 +71,39 @@ def _convert_extensions(operation: Dict[str, Any]) -> None:
                 "x-python-custom-pager-async", "azure.core.async_paging.AsyncItemPaged"
             )
         if operation_extensions.get("x-ms-long-running-operation"):
+            # poller
             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"
             )
 
+            # polling methods
+            sync_polling_method_directive = "x-python-custom-default-polling-method-sync"
+            operation["extensions"]["default-polling-method-sync"] = {
+                "azure-arm": operation_extensions.get(
+                    sync_polling_method_directive, "azure.mgmt.core.polling.arm_polling.ARMPolling"
+                ),
+                "data-plane": operation_extensions.get(
+                    sync_polling_method_directive, "azure.core.polling.base_polling.LROBasePolling"
+                ),
+            }
+            async_polling_method_directive = "x-python-custom-default-polling-method-async"
+            operation["extensions"]["default-polling-method-async"] = {
+                "azure-arm": operation_extensions.get(
+                    async_polling_method_directive, "azure.mgmt.core.polling.async_arm_polling.AsyncARMPolling"
+                ),
+                "data-plane": operation_extensions.get(
+                    async_polling_method_directive, "azure.core.polling.async_base_polling.AsyncLROBasePolling"
+                ),
+            }
+
+            operation["extensions"]["default-no-polling-method-sync"] = "azure.core.polling.NoPolling"
+            operation["extensions"]["default-no-polling-method-async"] = "azure.core.polling.AsyncNoPolling"
+            operation["extensions"]["base-polling-method-sync"] = "azure.core.polling.PollingMethod"
+            operation["extensions"]["base-polling-method-async"] = "azure.core.polling.AsyncPollingMethod"
+
     @staticmethod
     def _convert_schemas(schemas: Dict[str, Any]) -> None:
         for enum in schemas.get("sealedChoices", []) + schemas.get("choices", []):

test/azure/Expected/AcceptanceTests/CustomPollerPager/custompollerpager/aio/operations/_paging_operations.py

@@ -1124,8 +1124,8 @@ async def begin_get_multiple_pages_lro(
         :type paging_get_multiple_pages_lro_options: ~custompollerpager.models.PagingGetMultiplePagesLroOptions
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the AsyncARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.AsyncPollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of AsyncCustomPoller that returns an iterator like instance of either ProductResult or the result of cls(response)

test/azure/Expected/AcceptanceTests/CustomPollerPager/custompollerpager/operations/_paging_operations.py

@@ -1145,8 +1145,8 @@ def begin_get_multiple_pages_lro(
         :type paging_get_multiple_pages_lro_options: ~custompollerpager.models.PagingGetMultiplePagesLroOptions
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the ARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.PollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of CustomPoller that returns an iterator like instance of either ProductResult or the result of cls(response)

test/azure/Expected/AcceptanceTests/Lro/lro/aio/operations/_lr_os_custom_header_operations.py

@@ -108,8 +108,8 @@ async def begin_put_async_retry_succeeded(
         :type product: ~lro.models.Product
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the AsyncARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.AsyncPollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of AsyncLROPoller that returns either Product or the result of cls(response)
@@ -224,8 +224,8 @@ async def begin_put201_creating_succeeded200(
         :type product: ~lro.models.Product
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the AsyncARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.AsyncPollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of AsyncLROPoller that returns either Product or the result of cls(response)
@@ -331,8 +331,8 @@ async def begin_post202_retry200(
         :type product: ~lro.models.Product
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the AsyncARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.AsyncPollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of AsyncLROPoller that returns either None or the result of cls(response)
@@ -437,8 +437,8 @@ async def begin_post_async_retry_succeeded(
         :type product: ~lro.models.Product
         :keyword callable cls: A custom type or function that will be passed the direct response
         :keyword str continuation_token: A continuation token to restart a poller from a saved state.
-        :keyword polling: True for ARMPolling, False for no polling, or a
-         polling object for personal polling strategy
+        :keyword polling: Pass in True if you'd like the AsyncARMPolling polling method,
+         False for no polling, or your own initialized polling object for a personal polling strategy.
         :paramtype polling: bool or ~azure.core.polling.AsyncPollingMethod
         :keyword int polling_interval: Default waiting time between two polls for LRO operations if no Retry-After header is present.
         :return: An instance of AsyncLROPoller that returns either None or the result of cls(response)