fix content type logic

autorest/codegen/__init__.py

@@ -31,9 +31,9 @@ def _build_convenience_layer(yaml_data: Dict[str, Any], code_model: CodeModel) -
             code_model.sort_schemas()
 
     if code_model.options["show_operations"]:
-        code_model.link_operation_to_request_builder()
         code_model.add_schema_link_to_operation()
-        code_model.generate_single_parameter_from_multiple_media_types_operation()
+        code_model.generate_single_parameter_from_multiple_content_types_operation()
+        code_model.link_operation_to_request_builder()
         # LRO operation
         code_model.format_lro_operations()
         code_model.remove_next_operation()

autorest/codegen/models/__init__.py

@@ -10,7 +10,7 @@
 from .object_schema import ObjectSchema, get_object_schema, HiddenModelObjectSchema
 from .dictionary_schema import DictionarySchema
 from .list_schema import ListSchema
-from .primitive_schemas import get_primitive_schema, AnySchema, PrimitiveSchema
+from .primitive_schemas import get_primitive_schema, AnySchema, PrimitiveSchema, IOSchema
 from .enum_schema import EnumSchema, HiddenModelEnumSchema, get_enum_schema
 from .base_schema import BaseSchema
 from .constant_schema import ConstantSchema
@@ -62,6 +62,7 @@
     "RequestBuilderParameter",
     "HiddenModelObjectSchema",
     "ParameterStyle",
+    "IOSchema",
 ]
 
 def _generate_as_object_schema(yaml_data: Dict[str, Any]) -> bool:

autorest/codegen/models/base_builder.py

@@ -6,18 +6,24 @@
 from typing import Any, Callable, Dict, List, Optional, Union, TYPE_CHECKING
 from .base_model import BaseModel
 from .schema_response import SchemaResponse
+from .schema_request import SchemaRequest
 
 if TYPE_CHECKING:
     from . import ParameterListType
 
 
 _M4_HEADER_PARAMETERS = ["content_type", "accept"]
 
-def create_parameters(yaml_data: Dict[str, Any], code_model, parameter_creator: Callable):
+def create_parameters(
+    yaml_data: Dict[str, Any], code_model, parameter_creator: Callable
+):
     multiple_requests = len(yaml_data["requests"]) > 1
 
-    multiple_media_type_parameters = []
-    parameters = [parameter_creator(yaml, code_model=code_model) for yaml in yaml_data.get("parameters", [])]
+    multiple_content_type_parameters = []
+    parameters = [
+        parameter_creator(yaml, code_model=code_model)
+        for yaml in yaml_data.get("parameters", [])
+    ]
 
     for request in yaml_data["requests"]:
         for yaml in request.get("parameters", []):
@@ -26,14 +32,14 @@ def create_parameters(yaml_data: Dict[str, Any], code_model, parameter_creator:
             if name in _M4_HEADER_PARAMETERS:
                 parameters.append(parameter)
             elif multiple_requests:
-                parameter.has_multiple_media_types = True
-                multiple_media_type_parameters.append(parameter)
+                parameter.has_multiple_content_types = True
+                multiple_content_type_parameters.append(parameter)
             else:
                 parameters.append(parameter)
 
-    if multiple_media_type_parameters:
+    if multiple_content_type_parameters:
         body_parameters_name_set = set(
-            p.serialized_name for p in multiple_media_type_parameters
+            p.serialized_name for p in multiple_content_type_parameters
         )
         if len(body_parameters_name_set) > 1:
             raise ValueError(
@@ -53,7 +59,7 @@ def create_parameters(yaml_data: Dict[str, Any], code_model, parameter_creator:
         if parameter_original_id in parameters_index:
             parameter.original_parameter = parameters_index[parameter_original_id]
 
-    return parameters, multiple_media_type_parameters
+    return parameters, multiple_content_type_parameters
 
 class BaseBuilder(BaseModel):
     """Base class for Operations and Request Builders"""
@@ -65,6 +71,7 @@ def __init__(
         name: str,
         description: str,
         parameters: "ParameterListType",
+        schema_requests: List[SchemaRequest],
         responses: Optional[List[SchemaResponse]] = None,
         summary: Optional[str] = None,
     ) -> None:
@@ -75,6 +82,7 @@ def __init__(
         self.parameters = parameters
         self.responses = responses or []
         self.summary = summary
+        self.schema_requests = schema_requests
 
     @property
     def default_content_type_declaration(self) -> str:

autorest/codegen/models/code_model.py

@@ -316,7 +316,7 @@ def add_schema_link_to_operation(self) -> None:
             for operation in operation_group.operations:
                 for obj in chain(
                     operation.parameters,
-                    operation.multiple_media_type_parameters or [],
+                    operation.multiple_content_type_parameters or [],
                     operation.responses,
                     operation.exceptions,
                     chain.from_iterable(response.headers for response in operation.responses),
@@ -337,11 +337,11 @@ def add_schema_link_to_global_parameters(self) -> None:
         for parameter in self.global_parameters:
             self._populate_schema(parameter)
 
-    def generate_single_parameter_from_multiple_media_types_operation(self) -> None:
+    def generate_single_parameter_from_multiple_content_types_operation(self) -> None:
         for operation_group in self.operation_groups:
             for operation in operation_group.operations:
-                if operation.multiple_media_type_parameters:
-                    operation.convert_multiple_media_type_parameters()
+                if operation.multiple_content_type_parameters:
+                    operation.convert_multiple_content_type_parameters()
 
     @property
     def need_vendored_code(self) -> bool:
@@ -383,3 +383,4 @@ def link_operation_to_request_builder(self) -> None:
                 if isinstance(operation, LROOperation):
                     request_builder.name = request_builder.name + "_initial"
                 operation.request_builder = request_builder
+                operation.link_body_kwargs_to_body_params()

autorest/codegen/models/imports.py

@@ -95,3 +95,4 @@ def merge(self, file_import: "FileImport") -> None:
                 for package_name, module_list in package_list.items():
                     for module_name in module_list:
                         self._add_import(package_name, import_type, module_name, typing_section)
+        self.type_definitions.update(file_import.type_definitions)

autorest/codegen/models/lro_operation.py

@@ -11,6 +11,7 @@
 from .schema_response import SchemaResponse
 from .imports import ImportType, TypingSection
 from .base_schema import BaseSchema
+from .schema_request import SchemaRequest
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -24,7 +25,8 @@ def __init__(
         description: str,
         api_versions: Set[str],
         parameters: ParameterList,
-        multiple_media_type_parameters: ParameterList,
+        multiple_content_type_parameters: ParameterList,
+        schema_requests: List[SchemaRequest],
         summary: Optional[str] = None,
         responses: Optional[List[SchemaResponse]] = None,
         exceptions: Optional[List[SchemaResponse]] = None,
@@ -38,7 +40,8 @@ def __init__(
             description,
             api_versions,
             parameters,
-            multiple_media_type_parameters,
+            multiple_content_type_parameters,
+            schema_requests,
             summary,
             responses,
             exceptions,
@@ -86,7 +89,8 @@ def initial_operation(self) -> Operation:
             description="",
             api_versions=self.api_versions,
             parameters=self.parameters,
-            multiple_media_type_parameters=self.multiple_media_type_parameters,
+            schema_requests=self.schema_requests,
+            multiple_content_type_parameters=self.multiple_content_type_parameters,
             summary=self.summary,
             responses=self.responses,
             want_description_docstring=False,

autorest/codegen/models/object_schema.py

@@ -224,15 +224,15 @@ def serialization_type(self) -> str:
 
     @property
     def type_annotation(self) -> str:
-        return "Any"
+        return "JSONType"
 
     @property
     def operation_type_annotation(self) -> str:
-        return "Any"
+        return "JSONType"
 
     @property
     def docstring_type(self) -> str:
-        return "Any"
+        return "JSONType"
 
     @property
     def docstring_text(self) -> str:

autorest/codegen/models/operation.py

@@ -10,11 +10,13 @@
 from .base_builder import BaseBuilder, create_parameters
 from .imports import FileImport, ImportType, TypingSection
 from .schema_response import SchemaResponse
-from .parameter import get_parameter
+from .parameter import Parameter, get_parameter
 from .parameter_list import ParameterList, get_parameter_list
 from .base_schema import BaseSchema
 from .object_schema import ObjectSchema
 from .request_builder import RequestBuilder
+from .schema_request import SchemaRequest
+from .primitive_schemas import IOSchema
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -30,7 +32,8 @@ def __init__(
         description: str,
         api_versions: Set[str],
         parameters: ParameterList,
-        multiple_media_type_parameters: ParameterList,
+        multiple_content_type_parameters: ParameterList,
+        schema_requests: List[SchemaRequest],
         summary: Optional[str] = None,
         responses: Optional[List[SchemaResponse]] = None,
         exceptions: Optional[List[SchemaResponse]] = None,
@@ -44,11 +47,12 @@ def __init__(
             description=description,
             parameters=parameters,
             responses=responses,
+            schema_requests=schema_requests,
             summary=summary,
         )
-        self.multiple_media_type_parameters = multiple_media_type_parameters
+        self.multiple_content_type_parameters = multiple_content_type_parameters
         self.api_versions = api_versions
-        self.multiple_media_type_parameters = multiple_media_type_parameters
+        self.multiple_content_type_parameters = multiple_content_type_parameters
         self.exceptions = exceptions or []
         self.want_description_docstring = want_description_docstring
         self.want_tracing = want_tracing
@@ -79,28 +83,7 @@ def is_stream_response(self) -> bool:
 
     @property
     def body_kwargs_to_pass_to_request_builder(self) -> List[str]:
-        kwargs = []
-        if not self.parameters.has_body:
-            return []
-        body_kwarg_names = self.request_builder.parameters.body_kwarg_names
-        for kwarg in body_kwarg_names:
-            if kwarg == "content":
-                if self.request_builder.is_stream:
-                    kwargs.append("content")
-            else:
-                kwargs.append(kwarg)
-        if not kwargs and not self.parameters.body[0].constant:
-            kwargs.append("content")
-        return kwargs
-
-    @property
-    def serialized_body_kwarg(self) -> str:
-        # body serialization can be passed to either "json" or "content"
-        if "json" in self.body_kwargs_to_pass_to_request_builder:
-            return "json"
-        if not self.request_builder.is_stream:
-            return "content"
-        raise ValueError("You should not be trying to serialize this body")
+        return [p.serialized_name for p in self.request_builder.body_kwargs_to_get]
 
     @property
     def has_optional_return_type(self) -> bool:
@@ -168,7 +151,7 @@ def _imports_shared(self, async_mode: bool) -> FileImport: # pylint: disable=unu
         for param in self.parameters.method:
             file_import.merge(param.imports())
 
-        for param in self.multiple_media_type_parameters:
+        for param in self.multiple_content_type_parameters:
             file_import.merge(param.imports())
 
         for response in self.responses:
@@ -239,30 +222,72 @@ def imports(self, async_mode: bool) -> FileImport:
             file_import.add_from_import(
                 f"{relative_path}_vendor", "_convert_request", ImportType.LOCAL
             )
+        if self.code_model.options["version_tolerant"] and (
+            self.parameters.has_body or
+            any(r for r in self.responses if r.has_body)
+        ):
+            file_import.define_mypy_type("JSONType", "Any")
         return file_import
 
-    def convert_multiple_media_type_parameters(self) -> None:
+    def _get_body_param_from_body_kwarg(self, body_kwarg: Parameter) -> Parameter:
+        # determine which of the body parameters returned from m4 corresponds to this body_kwarg
+        if not self.multiple_content_type_parameters.has_body:
+            return self.parameters.body[0]
+        if body_kwarg.serialized_name == "data":
+            return next(p for p in self.multiple_content_type_parameters.body if p.is_data_input)
+        if body_kwarg.serialized_name == "files":
+            return next(p for p in self.multiple_content_type_parameters.body if p.is_multipart)
+        if body_kwarg.serialized_name == "json":
+            # first check if there's any non-binary. In the case of multiple content types, there's
+            # usually one binary (for content), and one schema parameter (for json)
+            try:
+                return next(
+                    p for p in self.multiple_content_type_parameters.body
+                    if not isinstance(p.schema, IOSchema)
+                )
+            except StopIteration:
+                return next(p for p in self.multiple_content_type_parameters.body if p.is_json_parameter)
+        return self.multiple_content_type_parameters.body[0]
+
+    def link_body_kwargs_to_body_params(self) -> None:
+        if not self.parameters.has_body:
+            return
+        body_kwargs = [
+            p for p in self.request_builder.parameters.body
+            if p.content_types
+        ]
+        if len(body_kwargs) == 1:
+            self.parameters.body[0].body_kwargs = [body_kwargs[0]]
+            return
+        for body_kwarg in body_kwargs:
+            body_param = self._get_body_param_from_body_kwarg(body_kwarg)
+            body_param.body_kwargs.append(body_kwarg)
+
+    def convert_multiple_content_type_parameters(self) -> None:
         type_annot = ", ".join([
             param.schema.operation_type_annotation
-            for param in self.multiple_media_type_parameters
+            for param in self.multiple_content_type_parameters
         ])
         docstring_type = " or ".join([
-            param.schema.docstring_type for param in self.multiple_media_type_parameters
+            param.schema.docstring_type for param in self.multiple_content_type_parameters
         ])
         try:
             # get an optional param with object first. These params are the top choice
             # bc they have more info about how to serialize the body
             chosen_parameter = next(
-                p for p in self.multiple_media_type_parameters if not p.required and isinstance(p.schema, ObjectSchema)
+                p for p in self.multiple_content_type_parameters
+                if not p.required and isinstance(p.schema, ObjectSchema)
             )
         except StopIteration:  # pylint: disable=broad-except
             # otherwise, we get the first optional param, if that exists. If not, we just grab the first one
-            optional_parameters = [p for p in self.multiple_media_type_parameters if not p.required]
-            chosen_parameter = optional_parameters[0] if optional_parameters else self.multiple_media_type_parameters[0]
+            optional_parameters = [p for p in self.multiple_content_type_parameters if not p.required]
+            chosen_parameter = (
+                optional_parameters[0] if optional_parameters else self.multiple_content_type_parameters[0]
+            )
         if not chosen_parameter:
             raise ValueError("You are missing a parameter that has multiple media types")
-        chosen_parameter.multiple_media_types_type_annot = f"Union[{type_annot}]"
-        chosen_parameter.multiple_media_types_docstring_type = docstring_type
+        chosen_parameter.multiple_content_types_type_annot = f"Union[{type_annot}]"
+        chosen_parameter.multiple_content_types_docstring_type = docstring_type
         self.parameters.append(chosen_parameter)
 
     @classmethod
@@ -272,16 +297,22 @@ def from_yaml(cls, yaml_data: Dict[str, Any], *, code_model) -> "Operation":
 
         parameter_creator = get_parameter(code_model).from_yaml
         parameter_list_creator = get_parameter_list(code_model)
-        parameters, multiple_media_type_parameters = create_parameters(yaml_data, code_model, parameter_creator)
+        schema_requests = [SchemaRequest.from_yaml(yaml, code_model=code_model) for yaml in yaml_data["requests"]]
+        parameters, multiple_content_type_parameters = create_parameters(
+            yaml_data, code_model, parameter_creator
+        )
 
         return cls(
             code_model=code_model,
             yaml_data=yaml_data,
             name=name,
             description=yaml_data["language"]["python"]["description"],
             api_versions=set(value_dict["version"] for value_dict in yaml_data["apiVersions"]),
-            parameters=parameter_list_creator(code_model, parameters),
-            multiple_media_type_parameters=parameter_list_creator(code_model, multiple_media_type_parameters),
+            parameters=parameter_list_creator(code_model, parameters, schema_requests),
+            multiple_content_type_parameters=parameter_list_creator(
+                code_model, multiple_content_type_parameters, schema_requests
+            ),
+            schema_requests=schema_requests,
             summary=yaml_data["language"]["python"].get("summary"),
             responses=[SchemaResponse.from_yaml(yaml) for yaml in yaml_data.get("responses", [])],
             # Exception with no schema means default exception, we don't store them