LRO Continuation Token [need azure-core 1.6.0]

ChangeLog.md

@@ -1,8 +1,17 @@
 # Change Log
 
-### Unreleased
+### 2020-06-03 - 5.1.0-preview.1
 Modelerfour version: 4.13.351
 
+**Disclaimer**
+
+This version requires azure-core 1.6.0 and contains features and bugfixes 5.0.0-preview.8
+
+**Features**
+
+- Refactor async LRO poller with a AsyncLROPoller class + "begin_" prefix
+- Add continuation_token kwargs to LRO methods
+
 **Bug Fixes**
 - Corrected generation of the item name of paging response when extracting data  #648
 - Corrected return type typing annotation for operations that return an optional body  #656

autorest/codegen/models/__init__.py

@@ -13,7 +13,7 @@
 from .enum_schema import EnumSchema
 from .base_schema import BaseSchema
 from .constant_schema import ConstantSchema
-from .imports import FileImport, ImportType
+from .imports import FileImport, ImportType, TypingSection
 from .lro_operation import LROOperation
 from .paging_operation import PagingOperation
 from .parameter import Parameter
@@ -36,6 +36,7 @@
     "EnumSchema",
     "FileImport",
     "ImportType",
+    "TypingSection",
     "PrimitiveSchema",
     "LROOperation",
     "Operation",

autorest/codegen/models/imports.py

@@ -3,20 +3,20 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from enum import Enum, auto
+from enum import Enum
 from typing import Dict, Optional, Set
 
 
-class ImportType(Enum):
-    STDLIB = auto()
-    THIRDPARTY = auto()
-    AZURECORE = auto()
-    LOCAL = auto()
+class ImportType(str, Enum):
+    STDLIB = "stdlib"
+    THIRDPARTY = "thirdparty"
+    AZURECORE = "azurecore"
+    LOCAL = "local"
 
-class TypingSection(Enum):
-    REGULAR = auto()  # this import is always a typing import
-    CONDITIONAL = auto()  # is a typing import when we're dealing with files that py2 will use, else regular
-    TYPING = auto()  # never a typing import
+class TypingSection(str, Enum):
+    REGULAR = "regular"  # this import is always a typing import
+    CONDITIONAL = "conditional"  # is a typing import when we're dealing with files that py2 will use, else regular
+    TYPING = "typing"  # never a typing import
 
 
 class FileImport:

autorest/codegen/models/lro_operation.py

@@ -89,7 +89,7 @@ def imports(self, code_model, async_mode: bool) -> FileImport:
         file_import.add_from_import("typing", "Union", ImportType.STDLIB, TypingSection.CONDITIONAL)
         if async_mode:
             file_import.add_from_import("typing", "Optional", ImportType.STDLIB, TypingSection.CONDITIONAL)
-            file_import.add_from_import("azure.core.polling", "async_poller", ImportType.AZURECORE)
+            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']:

autorest/codegen/serializers/metadata_serializer.py

@@ -4,7 +4,8 @@
 # license information.
 # --------------------------------------------------------------------------
 import copy
-from typing import List, Optional, Set, Tuple
+import json
+from typing import List, Optional, Set, Tuple, Dict
 from jinja2 import Environment
 from ..models import (
     CodeModel,
@@ -13,16 +14,42 @@
     LROOperation,
     PagingOperation,
     CredentialSchema,
-    ParameterList
+    ParameterList,
+    TypingSection,
+    ImportType
 )
-from .import_serializer import FileImportSerializer
 
 def _correct_credential_parameter(global_parameters: ParameterList, async_mode: bool) -> None:
     credential_param = [
         gp for gp in global_parameters.parameters if isinstance(gp.schema, CredentialSchema)
     ][0]
     credential_param.schema = CredentialSchema(async_mode=async_mode)
 
+def _json_serialize_imports(
+    imports: Dict[TypingSection, Dict[ImportType, Dict[str, Set[Optional[str]]]]]
+):
+    if not imports:
+        return None
+
+    json_serialize_imports = {}
+    # need to make name_import set -> list to make the dictionary json serializable
+    # not using an OrderedDict since we're iterating through a set and the order there varies
+    # going to sort the list instead
+
+    for typing_section_key, typing_section_value in imports.items():
+        json_import_type_dictionary = {}
+        for import_type_key, import_type_value in typing_section_value.items():
+            json_package_name_dictionary = {}
+            for package_name, name_imports in import_type_value.items():
+                name_import_ordered_list = []
+                if name_imports:
+                    name_import_ordered_list = list(name_imports)
+                    name_import_ordered_list.sort()
+                json_package_name_dictionary[package_name] = name_import_ordered_list
+            json_import_type_dictionary[import_type_key] = json_package_name_dictionary
+        json_serialize_imports[typing_section_key] = json_import_type_dictionary
+    return json.dumps(json_serialize_imports)
+
 
 class MetadataSerializer:
     def __init__(self, code_model: CodeModel, env: Environment) -> None:
@@ -99,11 +126,11 @@ def _is_paging(operation):
             is_paging=_is_paging,
             str=str,
             sync_mixin_imports=(
-                FileImportSerializer(sync_mixin_imports, is_python_3_file=False)
+                _json_serialize_imports(sync_mixin_imports.imports)
                 if sync_mixin_imports else None
             ),
             async_mixin_imports=(
-                FileImportSerializer(async_mixin_imports, is_python_3_file=True)
+                _json_serialize_imports(async_mixin_imports.imports)
                 if async_mixin_imports else None
             )
         )

autorest/codegen/templates/lro_operation.py.jinja2

@@ -2,11 +2,11 @@
 {% import 'operation_tools.jinja2' as op_tools %}
 {% set trace_decorator = "@distributed_trace_async" if async_mode else "@distributed_trace" %}
 {% set async_prefix = "Async" if async_mode else "" %}
-{% set poller = "async_poller" if async_mode else "LROPoller" %}
-{% set operation_name = operation.python_name if async_mode else "begin_"+operation.python_name %}
+{% set poller = "AsyncLROPoller" if async_mode else "LROPoller" %}
+{% set operation_name = "begin_"+operation.python_name %}
 {% macro return_docstring() %}
-:return: {{ "" if async_mode else "An instance of LROPoller that returns either " }}{{ operation.responses[0].schema.docstring_text if operation.responses[0].has_body else "None"}}{{ "," if async_mode }} or the result of cls(response)
-:rtype: {{"" if async_mode else "~azure.core.polling.LROPoller["}}{{ operation.responses[0].schema.docstring_type if operation.responses[0].has_body else "None" }}{{ "" if async_mode else "]" }}{% endmacro %}
+:return: An instance of {{ "Async" if async_mode }}LROPoller that returns either {{ operation.responses[0].schema.docstring_text if operation.responses[0].has_body else "None"}} or the result of cls(response)
+:rtype: ~azure.core.polling.{{ "Async" if async_mode }}LROPoller[{{ operation.responses[0].schema.docstring_type if operation.responses[0].has_body else "None" }}]{% endmacro %}
 {% macro param_documentation_string(parameter) %}:param {{ parameter.serialized_name }}: {{ parameter.description }}{% endmacro %}
 {% macro response_headers(response) %}
 response_headers = {
@@ -20,7 +20,7 @@ response_headers = {
 {% if code_model.options['tracing'] %}
 {{ trace_decorator }}
 {% endif %}
-{% set return_type_wrapper = "" if async_mode else "LROPoller" %}
+{% set return_type_wrapper = "AsyncLROPoller" if async_mode else "LROPoller" %}
 {{ 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) }}
@@ -43,6 +43,7 @@ response_headers = {
     :type {{ parameter.serialized_name }}: {{ parameter.schema.docstring_type }}
 {% endfor %}
     :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
@@ -56,13 +57,15 @@ response_headers = {
         'polling_interval',
         self._config.polling_interval
     )
-    raw_result = {{ keywords.await }}self._{{ operation.name }}_initial(
-        {% for parameter in operation.parameters.method %}
-        {{ parameter.serialized_name }}={{ parameter.serialized_name }},
-        {% endfor %}
-        cls=lambda x,y,z: x,
-        **kwargs
-    )
+    cont_token = kwargs.pop('continuation_token', None)  # type: Optional[str]
+    if cont_token is None:
+        raw_result = {{ keywords.await }}self._{{ operation.name }}_initial(
+            {% for parameter in operation.parameters.method %}
+            {{ parameter.serialized_name }}={{ parameter.serialized_name }},
+            {% endfor %}
+            cls=lambda x,y,z: x,
+            **kwargs
+        )
 
     kwargs.pop('error_map', None)
     kwargs.pop('content_type', None)
@@ -88,5 +91,13 @@ response_headers = {
     {% endif %}
     elif polling is False: polling_method = {{ async_prefix }}NoPolling()
     else: polling_method = polling
-    return {{ keywords.await }}{{ poller }}(self._client, raw_result, get_long_running_output, polling_method)
+    if cont_token:
+        return {{ poller }}.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)
 {{ operation_name }}.metadata = {'url': '{{ operation.url }}'}  # type: ignore

autorest/codegen/templates/metadata.json.jinja2

@@ -50,18 +50,16 @@
     },
     "operation_mixins": {
         {% for operation in mixin_operations %}
-        {{ operation.name | tojson }} : {
+        {% set operation_name = "begin_" + operation.name if is_lro(operation) else operation.name %}
+        {{ operation_name | tojson }} : {
             "sync": {
-                {% set sync_operation_name = "begin_" + operation.name if is_lro(operation) else operation.name %}
                 {% set sync_return_type_wrapper = "LROPoller" if is_lro(operation) else ("ItemPaged" if is_paging(operation) else "") %}
-                "operation_name": {{ sync_operation_name | tojson }},
-                "signature": {{ op_tools.method_signature(operation, sync_operation_name, False, False, sync_return_type_wrapper) | tojson }}
+                "signature": {{ op_tools.method_signature(operation, operation_name, False, False, sync_return_type_wrapper) | tojson }}
             },
             "async": {
                 {% set coroutine = False if is_paging(operation) else True %}
-                {% set async_return_type_wrapper = "AsyncItemPaged" if is_paging(operation) else "" %}
-                "operation_name": {{ operation.name | tojson }},
-                "signature": {{ op_tools.method_signature(operation, operation.name, True, coroutine, async_return_type_wrapper) | tojson }},
+                {% set async_return_type_wrapper = "AsyncLROPoller" if is_lro(operation) else ("AsyncItemPaged" if is_paging(operation) else "") %}
+                "signature": {{ op_tools.method_signature(operation, operation_name, True, coroutine, async_return_type_wrapper) | tojson }},
                 "coroutine": {{ coroutine | tojson }}
             },
             "doc": {{ op_tools.operation_docstring(operation) | tojson }},

autorest/multiapi/__init__.py

@@ -11,7 +11,8 @@
 from collections import defaultdict
 from pathlib import Path
 from typing import Dict, List, Tuple, Optional, cast, Any
-from .multiapi_serializer import MultiAPISerializer
+from .serializers import MultiAPISerializer, FileImportSerializer
+from .models import FileImport
 from ..jsonrpc import AutorestAPI
 
 from .. import Plugin
@@ -103,8 +104,8 @@ def there_is_a_rt_that_contains_api_version(rt_dict, api_version):
     # Operations at client level
     versioned_dict.update(
         {
-            operation_metadata[sync_or_async]["operation_name"]: operation_metadata[sync_or_async]["available_apis"]
-            for operation_metadata in mixin_operations.values()
+            operation_name: operation_metadata[sync_or_async]["available_apis"]
+            for operation_name, operation_metadata in mixin_operations.items()
         }
     )
     for operation, api_versions_list in versioned_dict.items():
@@ -181,13 +182,11 @@ def _build_operation_mixin_meta(self, paths_to_versions: List[Path]) -> Dict[str
                 mixin_operations.setdefault(func_name, {}).setdefault('async', {})
                 mixin_operations[func_name]['sync'].update({
                     "signature": func['sync']['signature'],
-                    "operation_name": func['sync']['operation_name'],
                     "doc": func['doc'],
                     "call": func['call']
                 })
                 mixin_operations[func_name]['async'].update({
                     "signature": func['async']['signature'],
-                    "operation_name": func['async']['operation_name'],
                     "coroutine": func['async']['coroutine'],
                     "doc": func['doc'],
                     "call": func['call']
@@ -249,6 +248,20 @@ def _parse_package_name_input(self) -> str:
             self.output_package_name = self.input_package_name
         return module_name
 
+    def _merge_mixin_imports_across_versions(
+        self, paths_to_versions: List[Path], async_mode: bool
+    ) -> FileImport:
+        imports = FileImport()
+        imports_to_load = "async_imports" if async_mode else "sync_imports"
+        for version_path in paths_to_versions:
+            metadata_json = json.loads(self._autorestapi.read_file(version_path / "_metadata.json"))
+            if not metadata_json.get('operation_mixins'):
+                continue
+            current_version_imports = FileImport(json.loads(metadata_json[imports_to_load]))
+            imports.merge(current_version_imports)
+
+        return imports
+
     def process(self) -> bool:
         _LOGGER.info("Generating multiapi client")
         # If True, means the auto-profile will consider preview versions.
@@ -326,6 +339,14 @@ def process(self) -> bool:
             versioned_operations_dict, mixin_operations, last_api_version, preview_mode, async_mode=True
         )
 
+        sync_imports = self._merge_mixin_imports_across_versions(
+            paths_to_versions, async_mode=False
+        )
+
+        async_imports = self._merge_mixin_imports_across_versions(
+            paths_to_versions, async_mode=True
+        )
+
         conf = {
             "client_name": metadata_json["client"]["name"],
             "package_name": self.output_package_name,
@@ -342,8 +363,8 @@ def process(self) -> bool:
             ),
             "config": metadata_json["config"],
             "global_parameters": metadata_json["global_parameters"],
-            "sync_imports": metadata_json["sync_imports"],
-            "async_imports": metadata_json["async_imports"]
+            "sync_imports": str(FileImportSerializer(sync_imports, is_python_3_file=False)),
+            "async_imports": str(FileImportSerializer(async_imports, is_python_3_file=True))
         }
 
         multiapi_serializer = MultiAPISerializer(

autorest/multiapi/models/__init__.py

@@ -0,0 +1,13 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+from .imports import ImportType, FileImport, TypingSection
+
+__all__ = [
+    "ImportType",
+    "FileImport",
+    "TypingSection"
+]