@autorest/python 5.15.0 → 5.16.0

package/autorest/codegen/models/lro_operation.py

@@ -4,7 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 import logging
-from typing import Dict, List, Any, Optional, Set, cast
+from typing import Dict, List, Any, Optional, Set, cast, TYPE_CHECKING
 from .imports import FileImport
 from .operation import Operation
 from .parameter_list import ParameterList
@@ -12,6 +12,10 @@ from .schema_response import SchemaResponse
 from .imports import ImportType, TypingSection
 from .base_schema import BaseSchema
 from .schema_request import SchemaRequest
+from .request_builder import RequestBuilder
+
+if TYPE_CHECKING:
+    from .code_model import CodeModel
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -19,8 +23,9 @@ _LOGGER = logging.getLogger(__name__)
 class LROOperation(Operation):
     def __init__(
         self,
-        code_model,
         yaml_data: Dict[str, Any],
+        code_model: "CodeModel",
+        request_builder: RequestBuilder,
         name: str,
         description: str,
         api_versions: Set[str],
@@ -31,11 +36,14 @@ class LROOperation(Operation):
         responses: Optional[List[SchemaResponse]] = None,
         exceptions: Optional[List[SchemaResponse]] = None,
         want_description_docstring: bool = True,
-        want_tracing: bool = True
+        want_tracing: bool = True,
+        *,
+        abstract: bool = False,
     ) -> None:
-        super(LROOperation, self).__init__(
-            code_model,
+        super().__init__(
             yaml_data,
+            code_model,
+            request_builder,
             name,
             description,
             api_versions,
@@ -46,9 +54,12 @@ class LROOperation(Operation):
             responses,
             exceptions,
             want_description_docstring,
-            want_tracing,
+            want_tracing=want_tracing,
+            abstract=abstract,
+        )
+        self.lro_options = yaml_data.get("extensions", {}).get(
+            "x-ms-long-running-operation-options", {}
         )
-        self.lro_options = yaml_data.get("extensions", {}).get("x-ms-long-running-operation-options", {})
         self.name = "begin_" + self.name
 
     @property
@@ -68,12 +79,14 @@ class LROOperation(Operation):
                 schema_types = {r.schema for r in responses_with_bodies}
                 response_schema = cast(BaseSchema, response.schema).serialization_type
                 _LOGGER.warning(
-                    "Multiple schema types in responses: %s. Choosing: %s", schema_types, response_schema
+                    "Multiple schema types in responses: %s. Choosing: %s",
+                    schema_types,
+                    response_schema,
                 )
             except IndexError:
                 raise ValueError(
-                    f"Your swagger is invalid because you have multiple response schemas for LRO" +
-                    f" method {self.python_name} and none of them have a 200 status code."
+                    f"Your swagger is invalid because you have multiple response schemas for LRO"
+                    + f" method {self.python_name} and none of them have a 200 status code."
                 )
 
         elif num_response_schemas:
@@ -83,8 +96,9 @@ class LROOperation(Operation):
     @property
     def initial_operation(self) -> Operation:
         operation = Operation(
-            self.code_model,
             yaml_data={},
+            code_model=self.code_model,
+            request_builder=self.code_model.lookup_request_builder(id(self.yaml_data)),
             name=self.name[5:] + "_initial",
             description="",
             api_versions=self.api_versions,
@@ -119,10 +133,14 @@ class LROOperation(Operation):
         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)
+        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]
+        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)
@@ -140,23 +158,35 @@ class LROOperation(Operation):
         file_import = super().imports_for_multiapi(async_mode)
         poller_import_path = ".".join(self.get_poller_path(async_mode).split(".")[:-1])
         poller = self.get_poller(async_mode)
-        file_import.add_submodule_import(poller_import_path, poller, ImportType.AZURECORE, TypingSection.CONDITIONAL)
+        file_import.add_submodule_import(
+            poller_import_path, poller, ImportType.AZURECORE, TypingSection.CONDITIONAL
+        )
         return file_import
 
     def imports(self, async_mode: bool, is_python3_file: bool) -> FileImport:
         file_import = self._imports_base(async_mode, is_python3_file)
-        file_import.add_submodule_import("typing", "Union", ImportType.STDLIB, TypingSection.CONDITIONAL)
+        file_import.add_submodule_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_submodule_import(poller_import_path, poller, ImportType.AZURECORE)
+        file_import.add_submodule_import(
+            poller_import_path, poller, ImportType.AZURECORE
+        )
 
         default_polling_method_import_path = ".".join(
-            self.get_default_polling_method_path(async_mode, self.code_model.options['azure_arm']).split(".")[:-1]
+            self.get_default_polling_method_path(
+                async_mode, self.code_model.options["azure_arm"]
+            ).split(".")[:-1]
+        )
+        default_polling_method = self.get_default_polling_method(
+            async_mode, self.code_model.options["azure_arm"]
         )
-        default_polling_method = self.get_default_polling_method(async_mode, self.code_model.options['azure_arm'])
         file_import.add_submodule_import(
-            default_polling_method_import_path, default_polling_method, ImportType.AZURECORE
+            default_polling_method_import_path,
+            default_polling_method,
+            ImportType.AZURECORE,
         )
 
         default_no_polling_method_import_path = ".".join(
@@ -164,17 +194,23 @@ class LROOperation(Operation):
         )
         default_no_polling_method = self.get_default_no_polling_method(async_mode)
         file_import.add_submodule_import(
-            default_no_polling_method_import_path, default_no_polling_method, ImportType.AZURECORE
+            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_submodule_import(base_polling_method_import_path, base_polling_method, ImportType.AZURECORE)
+        file_import.add_submodule_import(
+            base_polling_method_import_path, base_polling_method, ImportType.AZURECORE
+        )
         file_import.add_submodule_import("typing", "cast", ImportType.STDLIB)
         if async_mode:
-            file_import.add_submodule_import("typing", "Optional", ImportType.STDLIB, TypingSection.CONDITIONAL)
+            file_import.add_submodule_import(
+                "typing", "Optional", ImportType.STDLIB, TypingSection.CONDITIONAL
+            )
         if self.code_model.options["tracing"] and self.want_tracing:
             file_import.add_submodule_import(
                 f"azure.core.tracing.decorator{'_async' if async_mode else ''}",

package/autorest/codegen/models/lro_paging_operation.py

@@ -7,8 +7,8 @@ from .imports import FileImport
 from .lro_operation import LROOperation
 from .paging_operation import PagingOperation
 
-class LROPagingOperation(PagingOperation, LROOperation):
 
+class LROPagingOperation(PagingOperation, LROOperation):
     def imports(self, async_mode: bool, is_python3_file: bool) -> FileImport:
         lro_imports = LROOperation.imports(self, async_mode, is_python3_file)
         paging_imports = PagingOperation.imports(self, async_mode, is_python3_file)
@@ -19,6 +19,5 @@ class LROPagingOperation(PagingOperation, LROOperation):
 
     @property
     def success_status_code(self):
-        """The list of all successfull status code.
-        """
+        """The list of all successfull status code."""
         return [200]

package/autorest/codegen/models/object_schema.py

@@ -3,11 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import Any, Dict, List, Optional, Union, Type
+from typing import Any, Dict, List, Optional, Union, Type, TYPE_CHECKING
 from .base_schema import BaseSchema
 from .dictionary_schema import DictionarySchema
 from .property import Property
-from .imports import FileImport, ImportType, TypingSection
+from .imports import FileImport, ImportModel, ImportType, TypingSection
+
+if TYPE_CHECKING:
+    from .code_model import CodeModel
 
 
 class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
@@ -20,19 +23,28 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
     """
 
     def __init__(
-        self, namespace: str, yaml_data: Dict[str, Any], name: str, description: str = "", **kwargs
+        self,
+        yaml_data: Dict[str, Any],
+        code_model: "CodeModel",
+        name: str,
+        description: str = "",
+        **kwargs,
     ) -> None:
-        super(ObjectSchema, self).__init__(namespace=namespace, yaml_data=yaml_data)
+        super().__init__(yaml_data=yaml_data, code_model=code_model)
         self.name = name
         self.description = description
         self.max_properties: Optional[int] = kwargs.pop("max_properties", None)
         self.min_properties: Optional[int] = kwargs.pop("min_properties", None)
         self.properties: List[Property] = kwargs.pop("properties", [])
         self.is_exception: bool = kwargs.pop("is_exception", False)
-        self.base_models: Union[List[int], List["ObjectSchema"]] = kwargs.pop("base_models", [])
+        self.base_models: Union[List[int], List["ObjectSchema"]] = kwargs.pop(
+            "base_models", []
+        )
         self.subtype_map: Optional[Dict[str, str]] = kwargs.pop("subtype_map", None)
         self.discriminator_name: Optional[str] = kwargs.pop("discriminator_name", None)
-        self.discriminator_value: Optional[str] = kwargs.pop("discriminator_value", None)
+        self.discriminator_value: Optional[str] = kwargs.pop(
+            "discriminator_value", None
+        )
         self._created_json_template_representation = False
 
     @property
@@ -45,7 +57,7 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
 
     @property
     def docstring_type(self) -> str:
-        return f"~{self.namespace}.models.{self.name}"
+        return f"~{self.code_model.namespace}.models.{self.name}"
 
     @property
     def docstring_text(self) -> str:
@@ -79,18 +91,21 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
         # don't add additional properties, because there's not really a concept of
         # additional properties in the template
         representation = {
-            f'"{prop.original_swagger_name}"': prop.get_json_template_representation(**kwargs)
+            f'"{prop.original_swagger_name}"': prop.get_json_template_representation(
+                **kwargs
+            )
             for prop in [
-                p for p in self.properties
+                p
+                for p in self.properties
                 if not (p.is_discriminator or p.name == "additional_properties")
             ]
         }
         try:
             # add discriminator prop if there is one
             discriminator = next(p for p in self.properties if p.is_discriminator)
-            representation[
-                discriminator.original_swagger_name
-            ] = self.discriminator_value or discriminator.original_swagger_name
+            representation[discriminator.original_swagger_name] = (
+                self.discriminator_value or discriminator.original_swagger_name
+            )
         except StopIteration:
             pass
 
@@ -99,20 +114,10 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
         self._created_json_template_representation = False
         return representation
 
-    def get_files_and_data_template_representation(self, **kwargs: Any) -> Any:
-        object_schema_names = kwargs.get("object_schema_names", [])
-        object_schema_names.append(self.name)  # do tis to avoid circular
-        kwargs["object_schema_names"] = object_schema_names
-        return {
-            "{}".format(
-                prop.original_swagger_name
-            ): prop.get_files_and_data_template_representation(**kwargs)
-            for prop in self.properties
-        }
-
-
     @classmethod
-    def from_yaml(cls, namespace: str, yaml_data: Dict[str, Any], **kwargs) -> "ObjectSchema":
+    def from_yaml(
+        cls, yaml_data: Dict[str, Any], code_model: "CodeModel"
+    ) -> "ObjectSchema":
         """Returns a ClassType from the dict object constructed from a yaml file.
 
         WARNING: This guy might create an infinite loop.
@@ -123,11 +128,13 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
         :returns: A ClassType.
         :rtype: ~autorest.models.schema.ClassType
         """
-        obj = cls(namespace, yaml_data, "", description="")
-        obj.fill_instance_from_yaml(namespace, yaml_data)
+        obj = cls(yaml_data, code_model, "", description="")
+        obj.fill_instance_from_yaml(yaml_data, code_model)
         return obj
 
-    def fill_instance_from_yaml(self, namespace: str, yaml_data: Dict[str, Any], **kwargs) -> None:
+    def fill_instance_from_yaml(
+        self, yaml_data: Dict[str, Any], code_model: "CodeModel"
+    ) -> None:
         properties = []
         base_models = []
 
@@ -141,20 +148,24 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
                 for immediate_parent in immediate_parents:
                     if immediate_parent["type"] == "dictionary":
                         additional_properties_schema = DictionarySchema.from_yaml(
-                            namespace=namespace, yaml_data=immediate_parent, **kwargs
+                            yaml_data=immediate_parent, code_model=code_model
                         )
                         properties.append(
                             Property(
+                                yaml_data={},
+                                code_model=code_model,
                                 name="additional_properties",
                                 schema=additional_properties_schema,
                                 original_swagger_name="",
-                                yaml_data={},
-                                description="Unmatched properties from the message are deserialized to this collection."
+                                description=(
+                                    "Unmatched properties from the message are "
+                                    "deserialized to this collection."
+                                ),
                             )
                         )
                     elif (
-                        immediate_parent["language"]["default"]["name"] != name and
-                        immediate_parent['type'] == "object"
+                        immediate_parent["language"]["default"]["name"] != name
+                        and immediate_parent["type"] == "object"
                     ):
                         base_models.append(id(immediate_parent))
 
@@ -164,21 +175,22 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
             subtype_map = {}
             # map of discriminator value to child's name
             for children_yaml in yaml_data["discriminator"]["immediate"].values():
-                subtype_map[children_yaml["discriminatorValue"]] = children_yaml["language"]["python"]["name"]
+                subtype_map[children_yaml["discriminatorValue"]] = children_yaml[
+                    "language"
+                ]["python"]["name"]
         if yaml_data.get("properties"):
             properties += [
-                Property.from_yaml(p, has_additional_properties=len(properties) > 0, **kwargs)
+                Property.from_yaml(
+                    p, code_model, has_additional_properties=len(properties) > 0
+                )
                 for p in yaml_data["properties"]
             ]
         # this is to ensure that the attribute map type and property type are generated correctly
 
-
-
         description = yaml_data["language"]["python"]["description"]
         is_exception = False
-        exceptions_set = kwargs.pop("exceptions_set", None)
-        if exceptions_set:
-            if id(yaml_data) in exceptions_set:
+        if code_model.exception_ids:
+            if id(yaml_data) in code_model.exception_ids:
                 is_exception = True
 
         self.yaml_data = yaml_data
@@ -202,37 +214,50 @@ class ObjectSchema(BaseSchema):  # pylint: disable=too-many-instance-attributes
     @property
     def property_with_discriminator(self) -> Any:
         try:
-            return next(p for p in self.properties if getattr(p.schema, "discriminator_name", None))
+            return next(
+                p
+                for p in self.properties
+                if getattr(p.schema, "discriminator_name", None)
+            )
         except StopIteration:
             return None
 
     def imports(self) -> FileImport:
         file_import = FileImport()
         if self.is_exception:
-            file_import.add_submodule_import("azure.core.exceptions", "HttpResponseError", ImportType.AZURECORE)
+            file_import.add_submodule_import(
+                "azure.core.exceptions", "HttpResponseError", ImportType.AZURECORE
+            )
         return file_import
 
     def model_file_imports(self) -> FileImport:
         file_import = self.imports()
-        file_import.add_import("__init__", ImportType.LOCAL, typing_section=TypingSection.TYPING, alias="_models")
+        file_import.add_import(
+            "__init__",
+            ImportType.LOCAL,
+            typing_section=TypingSection.TYPING,
+            alias="_models",
+        )
         return file_import
 
-class HiddenModelObjectSchema(ObjectSchema):
 
+class HiddenModelObjectSchema(ObjectSchema):
     @property
     def serialization_type(self) -> str:
         return "object"
 
-    def type_annotation(self, *, is_operation_file: bool = False) -> str:  # pylint: disable=unused-argument
+    def type_annotation(
+        self, *, is_operation_file: bool = False  # pylint: disable=unused-argument
+    ) -> str:
         if self.xml_metadata:
             return "ET.Element"
-        return "JSONType"
+        return "JSON"
 
     @property
     def docstring_type(self) -> str:
         if self.xml_metadata:
             return "ET.Element"
-        return "JSONType"
+        return "JSON"
 
     @property
     def docstring_text(self) -> str:
@@ -242,11 +267,36 @@ class HiddenModelObjectSchema(ObjectSchema):
 
     def imports(self) -> FileImport:
         file_import = FileImport()
-        file_import.add_submodule_import("typing", "Any", ImportType.STDLIB, TypingSection.CONDITIONAL)
+        file_import.add_submodule_import(
+            "typing", "Any", ImportType.STDLIB, TypingSection.CONDITIONAL
+        )
         if self.xml_metadata:
-            file_import.add_submodule_import("xml.etree", "ElementTree", ImportType.STDLIB, alias="ET")
+            file_import.add_submodule_import(
+                "xml.etree", "ElementTree", ImportType.STDLIB, alias="ET"
+            )
+        file_import.add_import("sys", ImportType.STDLIB)
+        file_import.define_mypy_type(
+            "JSON",
+            "MutableMapping[str, Any] # pylint: disable=unsubscriptable-object",
+            None,
+            {
+                (3, 9): ImportModel(
+                    TypingSection.CONDITIONAL,
+                    ImportType.STDLIB,
+                    "collections.abc",
+                    submodule_name="MutableMapping",
+                ),
+                None: ImportModel(
+                    TypingSection.CONDITIONAL,
+                    ImportType.STDLIB,
+                    "typing",
+                    submodule_name="MutableMapping",
+                ),
+            },
+        )
         return file_import
 
+
 def get_object_schema(code_model) -> Type[ObjectSchema]:
     if code_model.options["models_mode"]:
         return ObjectSchema