@autorest/python 6.4.1 → 6.4.2

package/autorest/codegen/models/__init__.py

@@ -183,3 +183,4 @@ RequestBuilderType = Union[RequestBuilder, OverloadedRequestBuilder]
 ParameterType = Union[
     Parameter, RequestBuilderParameter, ClientParameter, ConfigParameter
 ]
+OperationType = Union[Operation, LROOperation, PagingOperation, LROPagingOperation]

package/autorest/codegen/models/base_builder.py

@@ -45,7 +45,6 @@ class BaseBuilder(
         parameters: ParameterListType,
         *,
         overloads=None,
-        want_tracing: bool = True,
     ) -> None:
         super().__init__(yaml_data=yaml_data, code_model=code_model)
         self.client = client
@@ -56,7 +55,7 @@ class BaseBuilder(
             overloads or []
         )
         self._summary: str = yaml_data.get("summary", "")
-        self.want_tracing = want_tracing
+        self.want_tracing: bool = yaml_data.get("wantTracing", True)
         self.group_name: str = yaml_data[
             "groupName"
         ]  # either operationGroup or client I am on

package/autorest/codegen/models/client.py

@@ -16,6 +16,8 @@ from .request_builder import (
     get_request_builder,
 )
 from .parameter import Parameter, ParameterMethodLocation
+from .lro_operation import LROOperation
+from .lro_paging_operation import LROPagingOperation
 
 ParameterListType = TypeVar(
     "ParameterListType",
@@ -24,6 +26,7 @@ ParameterListType = TypeVar(
 
 if TYPE_CHECKING:
     from .code_model import CodeModel
+    from . import OperationType
 
 
 class _ClientConfigBase(Generic[ParameterListType], BaseModel):
@@ -72,7 +75,7 @@ class Client(_ClientConfigBase[ClientGlobalParameterList]):
                 OperationGroup.from_yaml(op_group, code_model, self)
                 for op_group in self.yaml_data.get("operationGroups", [])
             ]
-            self.format_lro_operations()
+            self.link_lro_initial_operations()
 
     def _build_request_builders(
         self,
@@ -80,11 +83,21 @@ class Client(_ClientConfigBase[ClientGlobalParameterList]):
         request_builders: List[Union[RequestBuilder, OverloadedRequestBuilder]] = []
         for og_group in self.yaml_data["operationGroups"]:
             for operation_yaml in og_group["operations"]:
+                if operation_yaml["discriminator"] in ("lro", "lropaging"):
+                    continue
                 request_builder = get_request_builder(
                     operation_yaml,
                     code_model=self.code_model,
                     client=self,
                 )
+                if operation_yaml.get("isLroInitialOperation"):
+                    # we want to change the name
+                    request_builder.name = request_builder.get_name(
+                        request_builder.yaml_data["name"][1 : -len("_initial")],
+                        request_builder.yaml_data,
+                        request_builder.code_model,
+                        request_builder.client,
+                    )
                 if request_builder.overloads:
                     request_builders.extend(request_builder.overloads)
                 request_builders.append(request_builder)
@@ -160,6 +173,17 @@ class Client(_ClientConfigBase[ClientGlobalParameterList]):
                 f"No request builder with id {request_builder_id} found."
             ) from exc
 
+    def lookup_operation(self, operation_id: int) -> "OperationType":
+        try:
+            return next(
+                o
+                for og in self.operation_groups
+                for o in og.operations
+                if id(o.yaml_data) == operation_id
+            )
+        except StopIteration as exc:
+            raise KeyError(f"No operation with id {operation_id} found.") from exc
+
     def _imports_shared(self, async_mode: bool) -> FileImport:
         file_import = FileImport()
 
@@ -221,19 +245,14 @@ class Client(_ClientConfigBase[ClientGlobalParameterList]):
             for operation_group in self.operation_groups
         )
 
-    def format_lro_operations(self) -> None:
-        """Adds operations and attributes needed for LROs.
-        If there are LRO functions in here, will add initial LRO function. Will also set the return
-        type of the LRO operation
-        """
+    def link_lro_initial_operations(self) -> None:
+        """Link each LRO operation to its initial operation"""
         for operation_group in self.operation_groups:
-            i = 0
-            while i < len(operation_group.operations):
-                operation = operation_group.operations[i]
-                if operation.operation_type in ("lro", "lropaging"):
-                    operation_group.operations.insert(i, operation.initial_operation)  # type: ignore
-                    i += 1
-                i += 1
+            for operation in operation_group.operations:
+                if isinstance(operation, (LROOperation, LROPagingOperation)):
+                    operation.initial_operation = self.lookup_operation(
+                        id(operation.yaml_data["initialOperation"])
+                    )
 
     @property
     def need_request_converter(self) -> bool:

package/autorest/codegen/models/lro_operation.py

@@ -14,6 +14,7 @@ from .parameter_list import ParameterList
 if TYPE_CHECKING:
     from .code_model import CodeModel
     from .client import Client
+    from . import OperationType
 
 LROResponseType = TypeVar(
     "LROResponseType", bound=Union[LROResponse, LROPagingResponse]
@@ -33,8 +34,6 @@ class LROOperationBase(OperationBase[LROResponseType]):
         exceptions: List[Response],
         *,
         overloads: Optional[List[Operation]] = None,
-        public: bool = True,
-        want_tracing: bool = True,
     ) -> None:
         super().__init__(
             code_model=code_model,
@@ -46,11 +45,22 @@ class LROOperationBase(OperationBase[LROResponseType]):
             responses=responses,
             exceptions=exceptions,
             overloads=overloads,
-            public=public,
-            want_tracing=want_tracing,
         )
         self.name = "begin_" + self.name
         self.lro_options: Dict[str, Any] = self.yaml_data.get("lroOptions", {})
+        self._initial_operation: Optional["OperationType"] = None
+
+    @property
+    def initial_operation(self) -> "OperationType":
+        if not self._initial_operation:
+            raise ValueError(
+                "You need to first call client.link_lro_initial_operations before accessing"
+            )
+        return self._initial_operation
+
+    @initial_operation.setter
+    def initial_operation(self, val: "OperationType") -> None:
+        self._initial_operation = val
 
     @property
     def operation_type(self) -> str:
@@ -93,26 +103,6 @@ class LROOperationBase(OperationBase[LROResponseType]):
         """We don't want the poller to show up in ClsType, so we call super() on resposne type annotation"""
         return f"ClsType[{Response.type_annotation(self.responses[0], async_mode=async_mode)}]"
 
-    @property
-    def initial_operation(self) -> Operation:
-        """Initial operation that creates the first call for LRO polling"""
-        return Operation(
-            yaml_data=self.yaml_data,
-            code_model=self.code_model,
-            client=self.client,
-            request_builder=self.client.lookup_request_builder(id(self.yaml_data)),
-            name=self.name[5:] + "_initial",
-            overloads=self.overloads,
-            parameters=self.parameters,
-            responses=[
-                Response(r.yaml_data, self.code_model, headers=r.headers, type=r.type)
-                for r in self.responses
-            ],
-            exceptions=self.exceptions,
-            public=False,
-            want_tracing=False,
-        )
-
     def get_poller(self, async_mode: bool) -> str:
         return self.responses[0].get_poller(async_mode)
 
@@ -144,6 +134,10 @@ class LROOperationBase(OperationBase[LROResponseType]):
         file_import.add_submodule_import("typing", "cast", ImportType.STDLIB)
         return file_import
 
+    @classmethod
+    def get_request_builder(cls, yaml_data: Dict[str, Any], client: "Client"):
+        return client.lookup_request_builder(id(yaml_data["initialOperation"]))
+
 
 class LROOperation(LROOperationBase[LROResponse]):
     ...

package/autorest/codegen/models/operation.py

@@ -41,6 +41,7 @@ from .request_builder import OverloadedRequestBuilder, RequestBuilder
 if TYPE_CHECKING:
     from .code_model import CodeModel
     from .client import Client
+    from . import OperationType
 
 ResponseType = TypeVar(
     "ResponseType",
@@ -63,8 +64,6 @@ class OperationBase(  # pylint: disable=too-many-public-methods
         exceptions: List[Response],
         *,
         overloads: Optional[List["Operation"]] = None,
-        public: bool = True,
-        want_tracing: bool = True,
     ) -> None:
         super().__init__(
             code_model=code_model,
@@ -73,14 +72,16 @@ class OperationBase(  # pylint: disable=too-many-public-methods
             name=name,
             parameters=parameters,
             overloads=overloads,
-            want_tracing=want_tracing,
         )
         self.overloads: List["Operation"] = overloads or []
         self.responses = responses
-        self.public = public
         self.request_builder = request_builder
         self.deprecated = False
         self.exceptions = exceptions
+        self.is_lro_initial_operation: bool = self.yaml_data.get(
+            "isLroInitialOperation", False
+        )
+        self.include_documentation: bool = not self.is_lro_initial_operation
 
     @property
     def operation_type(self) -> str:
@@ -438,6 +439,10 @@ class OperationBase(  # pylint: disable=too-many-public-methods
     def has_stream_response(self) -> bool:
         return any(r.is_stream_response for r in self.responses)
 
+    @classmethod
+    def get_request_builder(cls, yaml_data: Dict[str, Any], client: "Client"):
+        return client.lookup_request_builder(id(yaml_data))
+
     @classmethod
     def from_yaml(
         cls,
@@ -446,7 +451,7 @@ class OperationBase(  # pylint: disable=too-many-public-methods
         client: "Client",
     ):
         name = yaml_data["name"]
-        request_builder = client.lookup_request_builder(id(yaml_data))
+        request_builder = cls.get_request_builder(yaml_data, client)
         responses = [
             cast(ResponseType, get_response(r, code_model))
             for r in yaml_data["responses"]
@@ -470,7 +475,6 @@ class OperationBase(  # pylint: disable=too-many-public-methods
             overloads=overloads,
             responses=responses,
             exceptions=exceptions,
-            want_tracing=not yaml_data["isOverload"],
         )
 
 
@@ -510,7 +514,7 @@ class Operation(OperationBase[Response]):
 
 def get_operation(
     yaml_data: Dict[str, Any], code_model: "CodeModel", client: "Client"
-) -> OperationBase:
+) -> "OperationType":
     if yaml_data["discriminator"] == "lropaging":
         from .lro_paging_operation import LROPagingOperation as OperationCls
     elif yaml_data["discriminator"] == "lro":

package/autorest/codegen/models/operation_group.py

@@ -8,13 +8,14 @@ from typing import Dict, List, Any, TYPE_CHECKING
 from autorest.codegen.models.utils import OrderedSet
 
 from .base import BaseModel
-from .operation import OperationBase, get_operation
+from .operation import get_operation
 from .imports import FileImport, ImportType, TypingSection
 from .utils import add_to_pylint_disable
 
 if TYPE_CHECKING:
     from .code_model import CodeModel
     from .client import Client
+    from . import OperationType
 
 
 class OperationGroup(BaseModel):
@@ -25,7 +26,7 @@ class OperationGroup(BaseModel):
         yaml_data: Dict[str, Any],
         code_model: "CodeModel",
         client: "Client",
-        operations: List[OperationBase],
+        operations: List["OperationType"],
         api_versions: List[str],
     ) -> None:
         super().__init__(yaml_data, code_model)

package/autorest/codegen/models/paging_operation.py

@@ -38,8 +38,6 @@ class PagingOperationBase(OperationBase[PagingResponseType]):
         exceptions: List[Response],
         *,
         overloads: Optional[List[Operation]] = None,
-        public: bool = True,
-        want_tracing: bool = True,
         override_success_response_to_200: bool = False,
     ) -> None:
         super().__init__(
@@ -52,8 +50,6 @@ class PagingOperationBase(OperationBase[PagingResponseType]):
             responses=responses,
             exceptions=exceptions,
             overloads=overloads,
-            public=public,
-            want_tracing=want_tracing,
         )
         self.next_request_builder: Optional[
             Union[RequestBuilder, OverloadedRequestBuilder]

package/autorest/codegen/models/request_builder.py

@@ -50,11 +50,11 @@ class RequestBuilderBase(BaseBuilder[ParameterListType]):
             name=name,
             parameters=parameters,
             overloads=overloads,
-            want_tracing=False,
         )
         self.overloads: List["RequestBuilder"] = overloads or []
         self.url: str = yaml_data["url"]
         self.method: str = yaml_data["method"]
+        self.want_tracing = False
 
     def response_type_annotation(self, **kwargs) -> str:
         return "HttpRequest"
@@ -127,14 +127,13 @@ class RequestBuilderBase(BaseBuilder[ParameterListType]):
         ...
 
     @classmethod
-    def from_yaml(
+    def get_name(
         cls,
+        name: str,
         yaml_data: Dict[str, Any],
         code_model: "CodeModel",
         client: "Client",
-    ):
-        # when combine embedded builders into one operation file, we need to avoid duplicated build function name.
-        # So add operation group name is effective method
+    ) -> str:
         additional_mark = ""
         if (
             code_model.options["combine_operation_files"]
@@ -146,10 +145,21 @@ class RequestBuilderBase(BaseBuilder[ParameterListType]):
         names = [
             "build",
             additional_mark,
-            yaml_data["name"],
+            name,
             "request",
         ]
-        name = "_".join([n for n in names if n])
+        return "_".join([n for n in names if n])
+
+    @classmethod
+    def from_yaml(
+        cls,
+        yaml_data: Dict[str, Any],
+        code_model: "CodeModel",
+        client: "Client",
+    ):
+        # when combine embedded builders into one operation file, we need to avoid duplicated build function name.
+        # So add operation group name is effective method
+
         overloads = [
             RequestBuilder.from_yaml(rb_yaml_data, code_model, client)
             for rb_yaml_data in yaml_data.get("overloads", [])
@@ -160,7 +170,7 @@ class RequestBuilderBase(BaseBuilder[ParameterListType]):
             yaml_data=yaml_data,
             code_model=code_model,
             client=client,
-            name=name,
+            name=cls.get_name(yaml_data["name"], yaml_data, code_model, client),
             parameters=parameter_list,
             overloads=overloads,
         )

package/autorest/codegen/serializers/metadata_serializer.py

@@ -5,7 +5,7 @@
 # --------------------------------------------------------------------------
 import functools
 import json
-from typing import List, Optional, Set, Tuple, Dict, Union
+from typing import List, Optional, Set, Tuple, Dict, Union, Any
 from jinja2 import Environment
 from ..models import (
     OperationGroup,
@@ -18,6 +18,12 @@ from ..models import (
 from .builder_serializer import get_operation_serializer
 
 
+def _to_string(data: Union[Tuple[Any], List[Any], str]) -> str:
+    if isinstance(data, (list, tuple)):
+        return "".join([_to_string(item) for item in data])
+    return str(data)
+
+
 def _json_serialize_imports(
     imports: Dict[
         TypingSection,
@@ -59,7 +65,7 @@ def _json_serialize_imports(
                 if name_imports:
                     name_import_ordered_list = list(name_imports)
                     name_import_ordered_list.sort(
-                        key=lambda e: "".join(e)  # type: ignore
+                        key=lambda e: _to_string(e)  # type: ignore
                         if isinstance(e, (list, tuple))
                         else e
                         if isinstance(e, str)

package/autorest/codegen/templates/lro_operation.py.jinja2

@@ -1,6 +1,6 @@
 {% import 'operation_tools.jinja2' as op_tools with context %}
 {# actual template starts here #}
-{% if operation.overloads and operation.public %}
+{% if operation.overloads and operation.include_documentation %}
 {{ op_tools.generate_overloads(operation_serializer, operation) }}
 {% endif %}
 {{ operation_serializer.method_signature_and_response_type_annotation(operation) }}

package/autorest/codegen/templates/lro_paging_operation.py.jinja2

@@ -1,7 +1,7 @@
 {% import 'operation_tools.jinja2' as op_tools with context %}
 {% import 'keywords.jinja2' as keywords with context %}
 {# actual template starts here #}
-{% if operation.overloads and operation.public %}
+{% if operation.overloads and operation.include_documentation %}
 {{ op_tools.generate_overloads(operation_serializer, operation) }}
 {% endif %}
 {{ operation_serializer.method_signature_and_response_type_annotation(operation) }}

package/autorest/codegen/templates/operation.py.jinja2

@@ -1,11 +1,11 @@
 {% import 'keywords.jinja2' as keywords with context %}
 {% import 'operation_tools.jinja2' as op_tools %}
 {# actual template starts here #}
-{% if operation.overloads and operation.public %}
+{% if operation.overloads and operation.include_documentation %}
 {{ op_tools.generate_overloads(operation_serializer, operation) }}
 {% endif %}
 {{ operation_serializer.method_signature_and_response_type_annotation(operation) }}
-{% if operation.public %}
+{% if operation.include_documentation %}
     {{ op_tools.description(operation, operation_serializer) | indent }}{% endif %}
     {% if not operation.abstract %}
         {% if operation.deprecated %}

package/autorest/codegen/templates/paging_operation.py.jinja2

@@ -1,10 +1,10 @@
 {% import 'operation_tools.jinja2' as op_tools with context %}
 {# actual template starts here #}
-{% if operation.overloads and operation.public %}
+{% if operation.overloads and operation.include_documentation %}
 {{ op_tools.generate_overloads(operation_serializer, operation) }}
 {% endif %}
 {{ operation_serializer.method_signature_and_response_type_annotation(operation) }}
-{% if operation.public %}
+{% if operation.include_documentation %}
     {{ op_tools.description(operation, operation_serializer) | indent }}{% endif %}
     {% if not operation.abstract %}
         {% if operation.deprecated %}

package/autorest/m4reformatter/__init__.py

@@ -299,10 +299,15 @@ def get_all_body_types(yaml_data: Dict[str, Any]) -> List[Dict[str, Any]]:
     return list(seen_body_types.values())
 
 
-def add_lro_information(operation: Dict[str, Any], yaml_data: Dict[str, Any]) -> None:
+def add_lro_information(
+    operation: Dict[str, Any],
+    initial_operation: Dict[str, Any],
+    yaml_data: Dict[str, Any],
+) -> None:
     operation["discriminator"] = "lro"
     extensions = yaml_data["extensions"]
     operation["lroOptions"] = extensions.get("x-ms-long-running-operation-options")
+    operation["initialOperation"] = initial_operation
     for response in operation["responses"]:
         response["pollerSync"] = extensions.get("x-python-custom-poller-sync")
         response["pollerAsync"] = extensions.get("x-python-custom-poller-async")
@@ -551,7 +556,7 @@ class M4Reformatter(
 
     def get_operation_creator(
         self, yaml_data: Dict[str, Any]
-    ) -> Callable[[str, Dict[str, Any]], Dict[str, Any]]:
+    ) -> Callable[[str, Dict[str, Any]], List[Dict[str, Any]]]:
         lro_operation = yaml_data.get("extensions", {}).get(
             "x-ms-long-running-operation"
         )
@@ -566,7 +571,7 @@ class M4Reformatter(
 
     def update_operation(
         self, group_name: str, yaml_data: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    ) -> List[Dict[str, Any]]:
         body_parameter = (
             self.update_body_parameter(yaml_data["requestMediaTypes"])
             if yaml_data.get("requestMediaTypes")
@@ -578,7 +583,7 @@ class M4Reformatter(
             group_name, yaml_data, body_parameter, content_types=content_types
         )
         operation["samples"] = yaml_data.get("extensions", {}).get("x-ms-examples", {})
-        return operation
+        return [operation]
 
     def add_paging_information(
         self, group_name: str, operation: Dict[str, Any], yaml_data: Dict[str, Any]
@@ -596,7 +601,7 @@ class M4Reformatter(
                 yaml_data=yaml_data["language"]["default"]["paging"][
                     "nextLinkOperation"
                 ],
-            )
+            )[0]
         extensions = yaml_data["extensions"]
         for response in operation["responses"]:
             response["pagerSync"] = extensions.get("x-python-custom-pager-sync")
@@ -604,27 +609,43 @@ class M4Reformatter(
 
     def update_paging_operation(
         self, group_name: str, yaml_data: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        base_operation = self.update_operation(group_name, yaml_data)
-        self.add_paging_information(group_name, base_operation, yaml_data)
-        return base_operation
+    ) -> List[Dict[str, Any]]:
+        retval: List[Dict[str, Any]] = []
+        for base_operation in self.update_operation(group_name, yaml_data):
+            self.add_paging_information(group_name, base_operation, yaml_data)
+            retval.append(base_operation)
+        return retval
 
     def update_lro_paging_operation(
         self, group_name: str, yaml_data: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        operation = self.update_lro_operation(group_name, yaml_data)
-        self.add_paging_information(group_name, operation, yaml_data)
-        operation["discriminator"] = "lropaging"
-        return operation
+    ) -> List[Dict[str, Any]]:
+        retval: List[Dict[str, Any]] = []
+        for operation in self.update_lro_operation(group_name, yaml_data):
+            if operation.get("discriminator") == "lro":
+                self.add_paging_information(group_name, operation, yaml_data)
+                operation["discriminator"] = "lropaging"
+            retval.append(operation)
+        return retval
+
+    def update_lro_initial_operation(self, initial_operation: Dict[str, Any]):
+        initial_operation["name"] = f"_{initial_operation['name']}_initial"
+        initial_operation["isLroInitialOperation"] = True
+        initial_operation["wantTracing"] = False
+        return initial_operation
 
     def update_lro_operation(
         self, group_name: str, yaml_data: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        base_operation = self.update_operation(group_name, yaml_data)
-        add_lro_information(base_operation, yaml_data)
-        for overload in base_operation["overloads"]:
-            add_lro_information(overload, yaml_data)
-        return base_operation
+    ) -> List[Dict[str, Any]]:
+        retval: List[Dict[str, Any]] = []
+        for base_operation in self.update_operation(group_name, yaml_data):
+            initial_operation = self.update_operation(group_name, yaml_data)[0]
+            self.update_lro_initial_operation(initial_operation)
+            add_lro_information(base_operation, initial_operation, yaml_data)
+            for overload in base_operation["overloads"]:
+                add_lro_information(overload, initial_operation, yaml_data)
+            retval.append(initial_operation)
+            retval.append(base_operation)
+        return retval
 
     def update_overload(
         self,
@@ -648,8 +669,9 @@ class M4Reformatter(
             "className": property_name,
             "operations": filter_out_paging_next_operation(
                 [
-                    self.get_operation_creator(o)(property_name, o)
-                    for o in yaml_data["operations"]
+                    o
+                    for ydo in yaml_data["operations"]
+                    for o in self.get_operation_creator(ydo)(property_name, ydo)
                 ]
             ),
         }

package/autorest/preprocess/__init__.py

@@ -68,6 +68,8 @@ def add_overload(
     overload["bodyParameter"]["type"] = body_type
     overload["bodyParameter"]["defaultToUnsetSentinel"] = False
     overload["overloads"] = []
+    if yaml_data.get("initialOperation"):
+        overload["initialOperation"] = yaml_data["initialOperation"]
 
     if for_flatten_params:
         overload["bodyParameter"]["flattened"] = True
@@ -320,9 +322,15 @@ class PreProcessPlugin(YamlUpdatePlugin):  # pylint: disable=abstract-method
         is_overload: bool = False,
     ) -> None:
         self.update_operation(code_model, yaml_data, is_overload=is_overload)
+        self.update_operation(
+            code_model, yaml_data["initialOperation"], is_overload=is_overload
+        )
         self._update_lro_operation_helper(yaml_data)
         for overload in yaml_data.get("overloads", []):
             self._update_lro_operation_helper(overload)
+            self.update_operation(
+                code_model, overload["initialOperation"], is_overload=True
+            )
 
     def update_paging_operation(
         self,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@autorest/python",
-  "version": "6.4.1",
+  "version": "6.4.2",
   "description": "The Python extension for generators in AutoRest.",
   "main": "index.js",
   "repository": {