@autorest/python 5.13.0 → 5.16.0

package/autorest/codegen/models/credential_schema.py

@@ -7,10 +7,14 @@ from typing import Any
 from .base_schema import BaseSchema
 from .imports import FileImport, ImportType, TypingSection
 
+
 class CredentialSchema(BaseSchema):
     def __init__(self) -> None:  # pylint: disable=super-init-not-called
         self.default_value = None
 
+    def type_annotation(self, *, is_operation_file: bool = False) -> str:
+        raise ValueError("Children classes should set their own type annotation")
+
     @property
     def docstring_type(self) -> str:
         return self.serialization_type
@@ -26,20 +30,19 @@ class CredentialSchema(BaseSchema):
         pass
 
     def get_json_template_representation(self, **kwargs: Any) -> Any:
-        raise TypeError("You should not try to get a JSON template representation of a CredentialSchema")
-
-    def get_files_and_data_template_representation(self, **kwargs: Any) -> Any:
-        raise TypeError("You should not try to get a files template representation of a CredentialSchema")
+        raise TypeError(
+            "You should not try to get a JSON template representation of a CredentialSchema"
+        )
 
 
 class AzureKeyCredentialSchema(CredentialSchema):
-
     @property
     def serialization_type(self) -> str:
         return "~azure.core.credentials.AzureKeyCredential"
 
-    @property
-    def type_annotation(self) -> str:
+    def type_annotation(
+        self, *, is_operation_file: bool = False  # pylint: disable=unused-argument
+    ) -> str:
         return "AzureKeyCredential"
 
     def imports(self) -> FileImport:
@@ -48,14 +51,14 @@ class AzureKeyCredentialSchema(CredentialSchema):
             "azure.core.credentials",
             "AzureKeyCredential",
             ImportType.AZURECORE,
-            typing_section=TypingSection.CONDITIONAL
+            typing_section=TypingSection.CONDITIONAL,
         )
         return file_import
 
 
 class TokenCredentialSchema(CredentialSchema):
     def __init__(self, async_mode) -> None:
-        super(TokenCredentialSchema, self).__init__()
+        super().__init__()
         self.async_mode = async_mode
         self.async_type = "~azure.core.credentials_async.AsyncTokenCredential"
         self.sync_type = "~azure.core.credentials.TokenCredential"
@@ -66,25 +69,27 @@ class TokenCredentialSchema(CredentialSchema):
             return self.async_type
         return self.sync_type
 
-    @property
-    def type_annotation(self) -> str:
+    def type_annotation(
+        self, *, is_operation_file: bool = False  # pylint: disable=unused-argument
+    ) -> str:
         if self.async_mode:
             return '"AsyncTokenCredential"'
         return '"TokenCredential"'
 
-
     def imports(self) -> FileImport:
         file_import = FileImport()
         if self.async_mode:
             file_import.add_submodule_import(
-                "azure.core.credentials_async", "AsyncTokenCredential",
+                "azure.core.credentials_async",
+                "AsyncTokenCredential",
                 ImportType.AZURECORE,
-                typing_section=TypingSection.TYPING
+                typing_section=TypingSection.TYPING,
             )
         else:
             file_import.add_submodule_import(
-                "azure.core.credentials", "TokenCredential",
+                "azure.core.credentials",
+                "TokenCredential",
                 ImportType.AZURECORE,
-                typing_section=TypingSection.TYPING
+                typing_section=TypingSection.TYPING,
             )
         return file_import

package/autorest/codegen/models/credential_schema_policy.py

@@ -7,9 +7,9 @@ from abc import abstractmethod
 from typing import List
 from .credential_schema import CredentialSchema
 
-class CredentialSchemaPolicy:
 
-    def __init__(self, credential: CredentialSchema, *args, **kwargs) -> None:  # pylint: disable=unused-argument
+class CredentialSchemaPolicy:
+    def __init__(self, credential: CredentialSchema) -> None:
         self.credential = credential
 
     @abstractmethod
@@ -22,11 +22,8 @@ class CredentialSchemaPolicy:
 
 
 class BearerTokenCredentialPolicy(CredentialSchemaPolicy):
-
     def __init__(
-        self,
-        credential: CredentialSchema,
-        credential_scopes: List[str]
+        self, credential: CredentialSchema, credential_scopes: List[str]
     ) -> None:
         super().__init__(credential)
         self._credential_scopes = credential_scopes
@@ -41,18 +38,14 @@ class BearerTokenCredentialPolicy(CredentialSchemaPolicy):
 
 
 class ARMChallengeAuthenticationPolicy(BearerTokenCredentialPolicy):
-
     def call(self, async_mode: bool) -> str:
         policy_name = f"Async{self.name()}" if async_mode else self.name()
         return f"{policy_name}(self.credential, *self.credential_scopes, **kwargs)"
 
 
 class AzureKeyCredentialPolicy(CredentialSchemaPolicy):
-
     def __init__(
-        self,
-        credential: CredentialSchema,
-        credential_key_header_name: str
+        self, credential: CredentialSchema, credential_key_header_name: str
     ) -> None:
         super().__init__(credential)
         self._credential_key_header_name = credential_key_header_name
@@ -64,14 +57,17 @@ class AzureKeyCredentialPolicy(CredentialSchemaPolicy):
     def call(self, async_mode: bool) -> str:
         return f'policies.AzureKeyCredentialPolicy(self.credential, "{self.credential_key_header_name}", **kwargs)'
 
+
 def get_credential_schema_policy_type(name):
-    policies = [ARMChallengeAuthenticationPolicy, BearerTokenCredentialPolicy, AzureKeyCredentialPolicy]
+    policies = [
+        ARMChallengeAuthenticationPolicy,
+        BearerTokenCredentialPolicy,
+        AzureKeyCredentialPolicy,
+    ]
     try:
         return next(p for p in policies if p.name().lower() == name.lower())
     except StopIteration:
         raise ValueError(
             "The credential policy you pass in with --credential-default-policy-type must be either "
-            "{}".format(
-                " or ".join([p.name() for p in policies])
-            )
+            "{}".format(" or ".join([p.name() for p in policies]))
         )

package/autorest/codegen/models/dictionary_schema.py

@@ -3,10 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, TYPE_CHECKING
 from .base_schema import BaseSchema
 from .imports import FileImport, ImportType, TypingSection
 
+if TYPE_CHECKING:
+    from .code_model import CodeModel
+
+
 class DictionarySchema(BaseSchema):
     """Schema for dictionaries that will be serialized.
 
@@ -18,11 +22,11 @@ class DictionarySchema(BaseSchema):
 
     def __init__(
         self,
-        namespace: str,
         yaml_data: Dict[str, Any],
-        element_type: "BaseSchema"
+        code_model: "CodeModel",
+        element_type: "BaseSchema",
     ) -> None:
-        super(DictionarySchema, self).__init__(namespace=namespace, yaml_data=yaml_data)
+        super().__init__(yaml_data=yaml_data, code_model=code_model)
         self.element_type = element_type
 
     @property
@@ -34,18 +38,13 @@ class DictionarySchema(BaseSchema):
         """
         return f"{{{self.element_type.serialization_type}}}"
 
-    @property
-    def type_annotation(self) -> str:
+    def type_annotation(self, *, is_operation_file: bool = False) -> str:
         """The python type used for type annotation
 
         :return: The type annotation for this schema
         :rtype: str
         """
-        return f"Dict[str, {self.element_type.type_annotation}]"
-
-    @property
-    def operation_type_annotation(self) -> str:
-        return f"Dict[str, {self.element_type.operation_type_annotation}]"
+        return f"Dict[str, {self.element_type.type_annotation(is_operation_file=is_operation_file)}]"
 
     @property
     def docstring_text(self) -> str:
@@ -60,20 +59,19 @@ class DictionarySchema(BaseSchema):
         return f"dict[str, {self.element_type.docstring_type}]"
 
     def xml_serialization_ctxt(self) -> Optional[str]:
-        raise NotImplementedError("Dictionary schema does not support XML serialization.")
+        raise NotImplementedError(
+            "Dictionary schema does not support XML serialization."
+        )
 
     def get_json_template_representation(self, **kwargs: Any) -> Any:
         return {
-            f'"{"str"}"' : self.element_type.get_json_template_representation(**kwargs)
-        }
-
-    def get_files_and_data_template_representation(self, **kwargs: Any) -> Any:
-        return {
-            f'"{"str"}"' : self.element_type.get_files_and_data_template_representation(**kwargs)
+            f'"{"str"}"': self.element_type.get_json_template_representation(**kwargs)
         }
 
     @classmethod
-    def from_yaml(cls, namespace: str, yaml_data: Dict[str, Any], **kwargs: Any) -> "DictionarySchema":
+    def from_yaml(
+        cls, yaml_data: Dict[str, Any], code_model: "CodeModel"
+    ) -> "DictionarySchema":
         """Constructs a DictionarySchema from yaml data.
 
         :param yaml_data: the yaml data from which we will construct this schema
@@ -86,18 +84,23 @@ class DictionarySchema(BaseSchema):
 
         from . import build_schema  # pylint: disable=import-outside-toplevel
 
-        element_type = build_schema(
-            yaml_data=element_schema, **kwargs
-        )
+        element_type = build_schema(yaml_data=element_schema, code_model=code_model)
 
         return cls(
-            namespace=namespace,
             yaml_data=yaml_data,
+            code_model=code_model,
             element_type=element_type,
         )
 
     def imports(self) -> FileImport:
         file_import = FileImport()
-        file_import.add_submodule_import("typing", "Dict", ImportType.STDLIB, TypingSection.CONDITIONAL)
+        file_import.add_submodule_import(
+            "typing", "Dict", ImportType.STDLIB, TypingSection.CONDITIONAL
+        )
         file_import.merge(self.element_type.imports())
         return file_import
+
+    def model_file_imports(self) -> FileImport:
+        file_import = self.imports()
+        file_import.merge(self.element_type.model_file_imports())
+        return file_import

package/autorest/codegen/models/enum_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, Set, Type
+from typing import Any, Dict, List, Optional, Set, Type, TYPE_CHECKING
 from .base_schema import BaseSchema
 from .primitive_schemas import PrimitiveSchema, get_primitive_schema, StringSchema
 from .imports import FileImport, ImportType, TypingSection
 
+if TYPE_CHECKING:
+    from .code_model import CodeModel
+
 
 class EnumValue:
     """Model containing necessary information for a single value of an enum.
@@ -17,7 +20,9 @@ class EnumValue:
     :param str description: Optional. The description for this enum value
     """
 
-    def __init__(self, name: str, value: str, description: Optional[str] = None) -> None:
+    def __init__(
+        self, name: str, value: str, description: Optional[str] = None
+    ) -> None:
         self.name = name
         self.value = value
         self.description = description
@@ -53,19 +58,17 @@ class EnumSchema(BaseSchema):
 
     def __init__(
         self,
-        namespace: str,
         yaml_data: Dict[str, Any],
+        code_model: "CodeModel",
         description: str,
         name: str,
         values: List["EnumValue"],
         enum_type: PrimitiveSchema,
-        enum_file_name: str
     ) -> None:
-        super(EnumSchema, self).__init__(namespace=namespace, yaml_data=yaml_data)
+        super().__init__(yaml_data=yaml_data, code_model=code_model)
         self.description = description
         self.name = name
         self.values = values
-        self.enum_file_name = enum_file_name
         self.enum_type = enum_type
 
     def __lt__(self, other):
@@ -80,23 +83,13 @@ class EnumSchema(BaseSchema):
         """
         return self.enum_type.serialization_type
 
-    @property
-    def type_annotation(self) -> str:
+    def type_annotation(self, *, is_operation_file: bool = False) -> str:
         """The python type used for type annotation
 
         :return: The type annotation for this schema
         :rtype: str
         """
-        return f'Union[{self.enum_type.type_annotation}, "{self.name}"]'
-
-    @property
-    def operation_type_annotation(self) -> str:
-        """The python type used for type annotation
-
-        :return: The type annotation for this schema
-        :rtype: str
-        """
-        return f'Union[{self.enum_type.type_annotation}, "_models.{self.name}"]'
+        return f'Union[{self.enum_type.type_annotation(is_operation_file=is_operation_file)}, "_models.{self.name}"]'
 
     def get_declaration(self, value: Any) -> str:
         return self.enum_type.get_declaration(value)
@@ -107,9 +100,8 @@ class EnumSchema(BaseSchema):
 
     @property
     def docstring_type(self) -> str:
-        """The python type used for RST syntax input and type annotation.
-        """
-        return f"{self.enum_type.type_annotation} or ~{self.namespace}.models.{self.name}"
+        """The python type used for RST syntax input and type annotation."""
+        return f"{self.enum_type.type_annotation()} or ~{self.code_model.namespace}.models.{self.name}"
 
     @staticmethod
     def _get_enum_values(yaml_data: List[Dict[str, Any]]) -> List["EnumValue"]:
@@ -133,21 +125,23 @@ class EnumSchema(BaseSchema):
 
     def _template_kwargs(self, **kwargs: Any) -> Any:
         if len(self.values) == 1 and not kwargs.get("default_value_declaration"):
-            kwargs['default_value_declaration'] = self.enum_type.get_declaration(self.values[0].value)
+            kwargs["default_value_declaration"] = self.enum_type.get_declaration(
+                self.values[0].value
+            )
         description = kwargs.pop("description", "")
         kwargs["description"] = description
         return kwargs
 
-
     def get_json_template_representation(self, **kwargs: Any) -> Any:
         # for better display effect, use the only value instead of var type
-        return self.enum_type.get_json_template_representation(**self._template_kwargs(**kwargs))
-
-    def get_files_and_data_template_representation(self, **kwargs: Any) -> Any:
-        return self.enum_type.get_files_and_data_template_representation(**self._template_kwargs(**kwargs))
+        return self.enum_type.get_json_template_representation(
+            **self._template_kwargs(**kwargs)
+        )
 
     @classmethod
-    def from_yaml(cls, namespace: str, yaml_data: Dict[str, Any], **kwargs: Any) -> "EnumSchema":
+    def from_yaml(
+        cls, yaml_data: Dict[str, Any], code_model: "CodeModel"
+    ) -> "EnumSchema":
         """Constructs an EnumSchema from yaml data.
 
         :param yaml_data: the yaml data from which we will construct this schema
@@ -160,63 +154,47 @@ class EnumSchema(BaseSchema):
 
         # choice type doesn't always exist. if there is no choiceType, we default to string
         if yaml_data.get("choiceType"):
-            enum_type = get_primitive_schema(namespace, yaml_data["choiceType"])
+            enum_type = get_primitive_schema(yaml_data["choiceType"], code_model)
         else:
-            enum_type = StringSchema(namespace, {"type": "str"})
+            enum_type = StringSchema({"type": "str"}, code_model)
         values = EnumSchema._get_enum_values(yaml_data["choices"])
-        code_model = kwargs.pop("code_model")
-
         return cls(
-            namespace=namespace,
             yaml_data=yaml_data,
+            code_model=code_model,
             description=yaml_data["language"]["python"]["description"],
             name=name,
             values=values,
             enum_type=enum_type,
-            enum_file_name=f"_{code_model.module_name}_enums"
         )
 
     def imports(self) -> FileImport:
         file_import = FileImport()
-        file_import.add_submodule_import("typing", "Union", ImportType.STDLIB, TypingSection.CONDITIONAL)
+        file_import.add_submodule_import(
+            "typing", "Union", ImportType.STDLIB, TypingSection.CONDITIONAL
+        )
         file_import.merge(self.enum_type.imports())
         return file_import
 
-    def model_file_imports(self) -> FileImport:
-        imports = self.imports()
-        # we import every enum since we can get extremely long imports
-        # if we import my name
-        imports.add_submodule_import("." + self.enum_file_name, "*", ImportType.LOCAL)
-        return imports
 
 class HiddenModelEnumSchema(EnumSchema):
-
     def imports(self) -> FileImport:
         file_import = FileImport()
         file_import.merge(self.enum_type.imports())
         return file_import
 
-    @property
-    def type_annotation(self) -> str:
-        """The python type used for type annotation
-
-        :return: The type annotation for this schema
-        :rtype: str
-        """
-        return self.enum_type.type_annotation
-
-    @property
-    def operation_type_annotation(self) -> str:
+    def type_annotation(self, *, is_operation_file: bool = False) -> str:
         """The python type used for type annotation
 
         :return: The type annotation for this schema
         :rtype: str
         """
-        return self.enum_type.type_annotation
+        return self.enum_type.type_annotation(is_operation_file=is_operation_file)
 
     @property
     def docstring_text(self) -> str:
-        return f"{self.enum_type.type_annotation}. {self.extra_description_information}"
+        return (
+            f"{self.enum_type.type_annotation()}. {self.extra_description_information}"
+        )
 
     @property
     def extra_description_information(self):
@@ -228,17 +206,18 @@ class HiddenModelEnumSchema(EnumSchema):
         if len(possible_values) == 2:
             possible_values_str = " or ".join(possible_values)
         else:
-            possible_values_str = ", ".join(
-                possible_values[: len(possible_values) - 1]
-            ) + f", and {possible_values[-1]}"
+            possible_values_str = (
+                ", ".join(possible_values[: len(possible_values) - 1])
+                + f", and {possible_values[-1]}"
+            )
 
-        return "Possible values are: {}.".format(possible_values_str)
+        return "Known values are: {}.".format(possible_values_str)
 
     @property
     def docstring_type(self) -> str:
-        """The python type used for RST syntax input and type annotation.
-        """
-        return self.enum_type.type_annotation
+        """The python type used for RST syntax input and type annotation."""
+        return self.enum_type.type_annotation()
+
 
 def get_enum_schema(code_model) -> Type[EnumSchema]:
     if code_model.options["models_mode"]:

package/autorest/codegen/models/imports.py

@@ -4,7 +4,8 @@
 # license information.
 # --------------------------------------------------------------------------
 from enum import Enum
-from typing import Dict, List, Optional, Tuple, Union, Set
+from typing import Dict, List, Optional, Tuple, Union, Set, Mapping
+
 
 class ImportType(str, Enum):
     STDLIB = "stdlib"
@@ -12,11 +13,13 @@ class ImportType(str, Enum):
     AZURECORE = "azurecore"
     LOCAL = "local"
 
+
 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 ImportModel:
     def __init__(
         self,
@@ -36,11 +39,11 @@ class ImportModel:
     def __eq__(self, other):
         try:
             return (
-                self.typing_section == other.typing_section and
-                self.import_type == other.import_type and
-                self.module_name == other.module_name and
-                self.submodule_name == other.submodule_name and
-                self.alias == other.alias
+                self.typing_section == other.typing_section
+                and self.import_type == other.import_type
+                and self.module_name == other.module_name
+                and self.submodule_name == other.submodule_name
+                and self.alias == other.alias
             )
         except AttributeError:
             return False
@@ -52,18 +55,32 @@ class ImportModel:
                 retval += hash(getattr(self, attr))
         return retval
 
-class FileImport:
+
+class TypeDefinition:
     def __init__(
         self,
-        imports: List[ImportModel] = None
-    ) -> None:
+        sync_definition: str,
+        async_definition: str,
+        version_imports: Mapping[Optional[Tuple[int, int]], ImportModel] = None,
+    ):
+        # version_imports: a map of "python version -> ImportModel".
+        #                  The python version is in form of (major, minor), for instance (3, 9) stands for py3.9.
+        #                  If the python version is None, it's a default ImportModel.
+        self.sync_definition = sync_definition
+        self.async_definition = async_definition
+        self.version_imports = version_imports
+
+
+class FileImport:
+    def __init__(self, imports: List[ImportModel] = None) -> None:
         self.imports = imports or []
         # has sync and async type definitions
-        self.type_definitions: Dict[str, Tuple[str, str]] = {}
+        self.type_definitions: Dict[str, TypeDefinition] = {}
 
     def _append_import(self, import_model: ImportModel) -> None:
         if not any(
-            i for i in self.imports
+            i
+            for i in self.imports
             if all(
                 getattr(i, attr) == getattr(import_model, attr)
                 for attr in dir(i)
@@ -72,7 +89,9 @@ class FileImport:
         ):
             self.imports.append(import_model)
 
-    def get_imports_from_section(self, typing_section: TypingSection) -> List[ImportModel]:
+    def get_imports_from_section(
+        self, typing_section: TypingSection
+    ) -> List[ImportModel]:
         return [i for i in self.imports if i.typing_section == typing_section]
 
     def add_submodule_import(
@@ -83,32 +102,44 @@ class FileImport:
         typing_section: TypingSection = TypingSection.REGULAR,
         alias: Optional[str] = None,
     ) -> None:
-        """Add an import to this import block.
-        """
-        self._append_import(ImportModel(
-            typing_section=typing_section,
-            import_type=import_type,
-            module_name=module_name,
-            submodule_name=submodule_name,
-            alias=alias,
-        ))
+        """Add an import to this import block."""
+        self._append_import(
+            ImportModel(
+                typing_section=typing_section,
+                import_type=import_type,
+                module_name=module_name,
+                submodule_name=submodule_name,
+                alias=alias,
+            )
+        )
 
     def add_import(
         self,
         module_name: str,
         import_type: ImportType,
-        typing_section: TypingSection = TypingSection.REGULAR
+        typing_section: TypingSection = TypingSection.REGULAR,
+        alias: Optional[str] = None,
     ) -> None:
         # Implementation detail: a regular import is just a "from" with no from
-        self._append_import(ImportModel(
-            typing_section=typing_section,
-            import_type=import_type,
-            module_name=module_name,
-        ))
+        self._append_import(
+            ImportModel(
+                typing_section=typing_section,
+                import_type=import_type,
+                module_name=module_name,
+                alias=alias,
+            )
+        )
 
-    def define_mypy_type(self, type_name: str, type_value: str, async_type_value: Optional[str] = None):
-        self.add_submodule_import("typing", "TypeVar", ImportType.STDLIB, TypingSection.CONDITIONAL)
-        self.type_definitions[type_name] = (type_value, async_type_value or type_value)
+    def define_mypy_type(
+        self,
+        type_name: str,
+        type_value: str,
+        async_type_value: Optional[str] = None,
+        version_imports: Mapping[Optional[Tuple[int, int]], ImportModel] = None,
+    ):
+        self.type_definitions[type_name] = TypeDefinition(
+            type_value, async_type_value or type_value, version_imports
+        )
 
     def merge(self, file_import: "FileImport") -> None:
         """Merge the given file import format."""
@@ -116,23 +147,23 @@ class FileImport:
             self._append_import(i)
         self.type_definitions.update(file_import.type_definitions)
 
-    def to_dict(self) -> Dict[
+    def to_dict(
+        self,
+    ) -> Dict[
         TypingSection,
-        Dict[ImportType, Dict[str, Set[Optional[Union[str, Tuple[str, str]]]]]]
+        Dict[ImportType, Dict[str, Set[Optional[Union[str, Tuple[str, str]]]]]],
     ]:
         retval: Dict[
             TypingSection,
-            Dict[ImportType, Dict[str, Set[Optional[Union[str, Tuple[str, str]]]]]]
+            Dict[ImportType, Dict[str, Set[Optional[Union[str, Tuple[str, str]]]]]],
         ] = dict()
         for i in self.imports:
             name_import: Optional[Union[str, Tuple[str, str]]] = None
             if i.submodule_name:
-                name_import = (i.submodule_name, i.alias) if i.alias else i.submodule_name
-            retval.setdefault(
-                    i.typing_section, dict()
-                ).setdefault(
-                    i.import_type, dict()
-                ).setdefault(
-                    i.module_name, set()
-                ).add(name_import)
+                name_import = (
+                    (i.submodule_name, i.alias) if i.alias else i.submodule_name
+                )
+            retval.setdefault(i.typing_section, dict()).setdefault(
+                i.import_type, dict()
+            ).setdefault(i.module_name, set()).add(name_import)
         return retval