fix default polling method listed in docstring

autorest/codegen/models/lro_operation.py

@@ -89,37 +89,62 @@ def has_optional_return_type(self) -> bool:
         return False
 
     def get_poller_path(self, async_mode: bool) -> str:
-        extension_name = "poller-async" if async_mode else "poller-sync"
+        extension_name = "poller-{}sync".format("a" if async_mode else "")
         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 get_default_polling_method_path(self, async_mode: bool, azure_arm: bool) -> str:
+        extension_name = "default-polling-method-{}sync".format("a" if async_mode else "")
+        arm_mode = "azure-arm" if azure_arm else "data-plane"
+        return self.yaml_data["extensions"][extension_name][arm_mode]
+
+    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:
+        extension_name = "default-no-polling-method-{}sync".format("a" if async_mode else "")
+        return self.yaml_data["extensions"][extension_name]
+
+    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:
+        extension_name = "base-polling-method-{}sync".format("a" if async_mode else "")
+        return self.yaml_data["extensions"][extension_name]
+    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/templates/lro_operation_helper.jinja2

@@ -26,8 +26,7 @@
 {% 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.
-{% set default_polling_method = "ARMPolling"if code_model.options['azure_arm'] else "LROBasePolling" %}
-:keyword polling: Pass in True if you'd like the {{ keywords.async_class }}{{ default_polling_method }} polling method,
+: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.
@@ -70,12 +69,8 @@
     }
 
     {% endif %}
-    {% if code_model.options['azure_arm'] %}
-    if polling is True: polling_method = {{ keywords.async_class }}ARMPolling(lro_delay{{ lro_options }}{{ path_format_arguments }},  **kwargs)
-    {% else %}
-    if polling is True: polling_method = {{ keywords.async_class }}LROBasePolling(lro_delay{{ lro_options }}{{ path_format_arguments }},  **kwargs)
-    {% endif %}
-    elif polling is False: polling_method = {{ keywords.async_class }}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/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/Paging/paging/aio/operations/_paging_operations.py

@@ -1140,7 +1140,7 @@ async def begin_get_multiple_pages_lro(
         :type paging_get_multiple_pages_lro_options: ~paging.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: Pass in True if you'd like the AsyncARMPolling polling method,
+        :keyword polling: Pass in True if you'd like the AsyncLROBasePolling 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.

test/azure/Expected/AcceptanceTests/Paging/paging/operations/_paging_operations.py

@@ -1160,7 +1160,7 @@ def begin_get_multiple_pages_lro(
         :type paging_get_multiple_pages_lro_options: ~paging.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: Pass in True if you'd like the ARMPolling polling method,
+        :keyword polling: Pass in True if you'd like the LROBasePolling 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.