Fix deserialization error for LRO which has discriminator

.chronus/changes/deserialization-fix-2024-4-24-16-48-41.md

@@ -0,0 +1,8 @@
+---
+changeKind: fix
+packages:
+  - "@autorest/python"
+  - "@azure-tools/typespec-python"
+---
+
+Fix deserialization error for lro when return type has discriminator and succeed in initial response

packages/autorest.python/autorest/codegen/serializers/builder_serializer.py

@@ -1223,7 +1223,7 @@ def _extract_data_callback(self, builder: PagingOperationType) -> List[str]:
             if isinstance(response.type, ModelType) and not response.type.internal:
                 deserialize_type = f'"{response.serialization_type}"'
                 pylint_disable = ""
-            deserialized = f"self._deserialize(\n    {deserialize_type},{pylint_disable}\n    pipeline_response\n)"
+            deserialized = f"self._deserialize(\n {deserialize_type},{pylint_disable}\n pipeline_response\n)"
             retval.append(f"    deserialized = {deserialized}")
         elif self.code_model.options["models_mode"] == "dpg":
             # we don't want to generate paging models for DPG
@@ -1318,6 +1318,8 @@ def initial_call(self, builder: LROOperationType) -> List[str]:
         retval.append("        params=_params,")
         retval.append("        **kwargs")
         retval.append("    )")
+        retval.append(f"    {'await ' if self.async_mode else ''}raw_result.http_response.read() # type: ignore")
+
         retval.append("kwargs.pop('error_map', None)")
         return retval
 

packages/autorest.python/autorest/preprocess/__init__.py

@@ -412,12 +412,18 @@ def update_lro_operation(
         yaml_data: Dict[str, Any],
         is_overload: bool = False,
     ) -> None:
+        def convert_initial_operation_response_type(data: Dict[str, Any]) -> None:
+            for response in data.get("responses", []):
+                response["type"] = KNOWN_TYPES["binary"]
+
         self.update_operation(code_model, yaml_data, is_overload=is_overload)
         self.update_operation(code_model, yaml_data["initialOperation"], is_overload=is_overload)
+        convert_initial_operation_response_type(yaml_data["initialOperation"])
         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)
+            convert_initial_operation_response_type(overload["initialOperation"])
 
     def update_paging_operation(
         self,
@@ -466,6 +472,7 @@ def update_operation_groups(self, code_model: Dict[str, Any], client: Dict[str,
     def update_yaml(self, yaml_data: Dict[str, Any]) -> None:
         """Convert in place the YAML str."""
         self.update_types(yaml_data["types"])
+        yaml_data["types"] += KNOWN_TYPES.values()
         for client in yaml_data["clients"]:
             self.update_client(client)
             self.update_operation_groups(yaml_data, client)

packages/autorest.python/samples/specification/directives/generated/azure/directives/sample/_operations/_operations.py

@@ -8,7 +8,7 @@
 # --------------------------------------------------------------------------
 from io import IOBase
 import sys
-from typing import Any, Callable, Dict, IO, Iterable, Optional, Type, TypeVar, Union, cast, overload
+from typing import Any, Callable, Dict, IO, Iterable, Iterator, Optional, Type, TypeVar, Union, cast, overload
 
 from my.library import CustomDefaultPollingMethod, CustomPager, CustomPoller
 
@@ -74,7 +74,9 @@ def build_polling_paging_example_basic_paging_request(**kwargs: Any) -> HttpRequ
 
 class PollingPagingExampleOperationsMixin(PollingPagingExampleMixinABC):
 
-    def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = None, **kwargs: Any) -> Optional[JSON]:
+    def _basic_polling_initial(
+        self, product: Optional[Union[JSON, IO[bytes]]] = None, **kwargs: Any
+    ) -> Iterator[bytes]:
         error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
@@ -87,7 +89,7 @@ def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = Non
         _params = kwargs.pop("params", {}) or {}
 
         content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None))
-        cls: ClsType[Optional[JSON]] = kwargs.pop("cls", None)
+        cls: ClsType[Iterator[bytes]] = kwargs.pop("cls", None)
 
         content_type = content_type or "application/json"
         _json = None
@@ -109,7 +111,7 @@ def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = Non
         )
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = True
         pipeline_response: PipelineResponse = self._client._pipeline.run(  # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -122,17 +124,16 @@ def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = Non
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             raise HttpResponseError(response=response)
 
-        deserialized = None
         if response.status_code == 200:
-            if response.content:
-                deserialized = response.json()
-            else:
-                deserialized = None
+            deserialized = response.iter_bytes()
+
+        if response.status_code == 204:
+            deserialized = response.iter_bytes()
 
         if cls:
-            return cls(pipeline_response, deserialized, {})  # type: ignore
+            return cls(pipeline_response, cast(Iterator[bytes], deserialized), {})  # type: ignore
 
-        return deserialized  # type: ignore
+        return cast(Iterator[bytes], deserialized)  # type: ignore
 
     @overload
     def begin_basic_polling(
@@ -245,6 +246,7 @@ def begin_basic_polling(
                 params=_params,
                 **kwargs
             )
+            raw_result.http_response.read()  # type: ignore
         kwargs.pop("error_map", None)
 
         def get_long_running_output(pipeline_response):

packages/autorest.python/samples/specification/directives/generated/azure/directives/sample/aio/_operations/_operations.py

@@ -8,7 +8,7 @@
 # --------------------------------------------------------------------------
 from io import IOBase
 import sys
-from typing import Any, AsyncIterable, Callable, Dict, IO, Optional, Type, TypeVar, Union, cast, overload
+from typing import Any, AsyncIterable, AsyncIterator, Callable, Dict, IO, Optional, Type, TypeVar, Union, cast, overload
 
 from my.library.aio import AsyncCustomDefaultPollingMethod, AsyncCustomPager, AsyncCustomPoller
 
@@ -47,7 +47,7 @@ class PollingPagingExampleOperationsMixin(PollingPagingExampleMixinABC):
 
     async def _basic_polling_initial(
         self, product: Optional[Union[JSON, IO[bytes]]] = None, **kwargs: Any
-    ) -> Optional[JSON]:
+    ) -> AsyncIterator[bytes]:
         error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
@@ -60,7 +60,7 @@ async def _basic_polling_initial(
         _params = kwargs.pop("params", {}) or {}
 
         content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None))
-        cls: ClsType[Optional[JSON]] = kwargs.pop("cls", None)
+        cls: ClsType[AsyncIterator[bytes]] = kwargs.pop("cls", None)
 
         content_type = content_type or "application/json"
         _json = None
@@ -82,7 +82,7 @@ async def _basic_polling_initial(
         )
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = True
         pipeline_response: PipelineResponse = await self._client._pipeline.run(  # type: ignore # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -95,17 +95,16 @@ async def _basic_polling_initial(
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             raise HttpResponseError(response=response)
 
-        deserialized = None
         if response.status_code == 200:
-            if response.content:
-                deserialized = response.json()
-            else:
-                deserialized = None
+            deserialized = response.iter_bytes()
+
+        if response.status_code == 204:
+            deserialized = response.iter_bytes()
 
         if cls:
-            return cls(pipeline_response, deserialized, {})  # type: ignore
+            return cls(pipeline_response, cast(AsyncIterator[bytes], deserialized), {})  # type: ignore
 
-        return deserialized  # type: ignore
+        return cast(AsyncIterator[bytes], deserialized)  # type: ignore
 
     @overload
     async def begin_basic_polling(
@@ -218,6 +217,7 @@ async def begin_basic_polling(
                 params=_params,
                 **kwargs
             )
+            await raw_result.http_response.read()  # type: ignore
         kwargs.pop("error_map", None)
 
         def get_long_running_output(pipeline_response):

packages/autorest.python/samples/specification/multiapi/generated/azure/multiapi/sample/_operations_mixin.py

@@ -10,7 +10,7 @@
 # --------------------------------------------------------------------------
 from ._serialization import Serializer, Deserializer
 from io import IOBase
-from typing import Any, IO, Iterable, Optional, Union
+from typing import Any, IO, Iterable, Iterator, Optional, Union
 
 from azure.core.paging import ItemPaged
 from azure.core.polling import LROPoller

packages/autorest.python/samples/specification/multiapi/generated/azure/multiapi/sample/aio/_operations_mixin.py

@@ -10,7 +10,7 @@
 # --------------------------------------------------------------------------
 from .._serialization import Serializer, Deserializer
 from io import IOBase
-from typing import Any, AsyncIterable, IO, Optional, Union
+from typing import Any, AsyncIterable, AsyncIterator, IO, Optional, Union
 
 from azure.core.async_paging import AsyncItemPaged
 from azure.core.polling import AsyncLROPoller

packages/autorest.python/samples/specification/multiapi/generated/azure/multiapi/sample/v1/_metadata.json

@@ -95,8 +95,8 @@
         "operation_group_one": "OperationGroupOneOperations"
     },
     "operation_mixins": {
-        "sync_imports": "{\"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"Iterable\"]}}, \"regular\": {\"stdlib\": {\"typing\": [\"IO\", \"Optional\", \"Union\"], \"io\": [\"IOBase\"]}, \"local\": {\".\": [[\"models\", \"_models\"]]}, \"sdkcore\": {\"azure.core.polling\": [\"LROPoller\"], \"azure.core.paging\": [\"ItemPaged\"]}}}",
-        "async_imports": "{\"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"AsyncIterable\"]}}, \"regular\": {\"stdlib\": {\"typing\": [\"IO\", \"Optional\", \"Union\"], \"io\": [\"IOBase\"]}, \"local\": {\"..\": [[\"models\", \"_models\"]]}, \"sdkcore\": {\"azure.core.polling\": [\"AsyncLROPoller\"], \"azure.core.async_paging\": [\"AsyncItemPaged\"]}}}",
+        "sync_imports": "{\"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"Iterable\"]}}, \"regular\": {\"stdlib\": {\"typing\": [\"IO\", \"Iterator\", \"Optional\", \"Union\"], \"io\": [\"IOBase\"]}, \"local\": {\".\": [[\"models\", \"_models\"]]}, \"sdkcore\": {\"azure.core.polling\": [\"LROPoller\"], \"azure.core.paging\": [\"ItemPaged\"]}}}",
+        "async_imports": "{\"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"AsyncIterable\"]}}, \"regular\": {\"stdlib\": {\"typing\": [\"AsyncIterator\", \"IO\", \"Optional\", \"Union\"], \"io\": [\"IOBase\"]}, \"local\": {\"..\": [[\"models\", \"_models\"]]}, \"sdkcore\": {\"azure.core.polling\": [\"AsyncLROPoller\"], \"azure.core.async_paging\": [\"AsyncItemPaged\"]}}}",
         "sync_mixin_typing_definitions": "",
         "async_mixin_typing_definitions": "",
         "operations": {
@@ -115,14 +115,14 @@
             },
             "_test_lro_initial" : {
                 "sync": {
-                    "signature": "def _test_lro_initial(\n    self,\n    product: Optional[Union[_models.Product, IO[bytes]]] = None,\n    **kwargs: Any\n) -\u003e Optional[_models.Product]:\n",
-                    "doc": "\"\"\"Put in whatever shape of Product you want, will return a Product with id equal to 100.\n\n:param product: Product to put. Is either a Product type or a IO[bytes] type. Default value is\n None.\n:type product: ~azure.multiapi.sample.v1.models.Product or IO[bytes]\n:return: Product or None or the result of cls(response)\n:rtype: ~azure.multiapi.sample.v1.models.Product or None\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
+                    "signature": "def _test_lro_initial(\n    self,\n    product: Optional[Union[_models.Product, IO[bytes]]] = None,\n    **kwargs: Any\n) -\u003e Iterator[bytes]:\n",
+                    "doc": "\"\"\"Put in whatever shape of Product you want, will return a Product with id equal to 100.\n\n:param product: Product to put. Is either a Product type or a IO[bytes] type. Default value is\n None.\n:type product: ~azure.multiapi.sample.v1.models.Product or IO[bytes]\n:return: Iterator[bytes] or the result of cls(response)\n:rtype: Iterator[bytes]\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
                     "call": "product, **kwargs"
                 },
                 "async": {
                     "coroutine": true,
-                    "signature": "async def _test_lro_initial(\n    self,\n    product: Optional[Union[_models.Product, IO[bytes]]] = None,\n    **kwargs: Any\n) -\u003e Optional[_models.Product]:\n",
-                    "doc": "\"\"\"Put in whatever shape of Product you want, will return a Product with id equal to 100.\n\n:param product: Product to put. Is either a Product type or a IO[bytes] type. Default value is\n None.\n:type product: ~azure.multiapi.sample.v1.models.Product or IO[bytes]\n:return: Product or None or the result of cls(response)\n:rtype: ~azure.multiapi.sample.v1.models.Product or None\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
+                    "signature": "async def _test_lro_initial(\n    self,\n    product: Optional[Union[_models.Product, IO[bytes]]] = None,\n    **kwargs: Any\n) -\u003e AsyncIterator[bytes]:\n",
+                    "doc": "\"\"\"Put in whatever shape of Product you want, will return a Product with id equal to 100.\n\n:param product: Product to put. Is either a Product type or a IO[bytes] type. Default value is\n None.\n:type product: ~azure.multiapi.sample.v1.models.Product or IO[bytes]\n:return: AsyncIterator[bytes] or the result of cls(response)\n:rtype: AsyncIterator[bytes]\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
                     "call": "product, **kwargs"
                 }
             },
@@ -141,14 +141,14 @@
             },
             "_test_lro_and_paging_initial" : {
                 "sync": {
-                    "signature": "def _test_lro_and_paging_initial(\n    self,\n    client_request_id: Optional[str] = None,\n    test_lro_and_paging_options: Optional[_models.TestLroAndPagingOptions] = None,\n    **kwargs: Any\n) -\u003e _models.PagingResult:\n",
-                    "doc": "\"\"\"A long-running paging operation that includes a nextLink that has 10 pages.\n\n:param client_request_id: Default value is None.\n:type client_request_id: str\n:param test_lro_and_paging_options: Parameter group. Default value is None.\n:type test_lro_and_paging_options: ~azure.multiapi.sample.v1.models.TestLroAndPagingOptions\n:return: PagingResult or the result of cls(response)\n:rtype: ~azure.multiapi.sample.v1.models.PagingResult\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
+                    "signature": "def _test_lro_and_paging_initial(\n    self,\n    client_request_id: Optional[str] = None,\n    test_lro_and_paging_options: Optional[_models.TestLroAndPagingOptions] = None,\n    **kwargs: Any\n) -\u003e Iterator[bytes]:\n",
+                    "doc": "\"\"\"A long-running paging operation that includes a nextLink that has 10 pages.\n\n:param client_request_id: Default value is None.\n:type client_request_id: str\n:param test_lro_and_paging_options: Parameter group. Default value is None.\n:type test_lro_and_paging_options: ~azure.multiapi.sample.v1.models.TestLroAndPagingOptions\n:return: Iterator[bytes] or the result of cls(response)\n:rtype: Iterator[bytes]\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
                     "call": "client_request_id, test_lro_and_paging_options, **kwargs"
                 },
                 "async": {
                     "coroutine": true,
-                    "signature": "async def _test_lro_and_paging_initial(\n    self,\n    client_request_id: Optional[str] = None,\n    test_lro_and_paging_options: Optional[_models.TestLroAndPagingOptions] = None,\n    **kwargs: Any\n) -\u003e _models.PagingResult:\n",
-                    "doc": "\"\"\"A long-running paging operation that includes a nextLink that has 10 pages.\n\n:param client_request_id: Default value is None.\n:type client_request_id: str\n:param test_lro_and_paging_options: Parameter group. Default value is None.\n:type test_lro_and_paging_options: ~azure.multiapi.sample.v1.models.TestLroAndPagingOptions\n:return: PagingResult or the result of cls(response)\n:rtype: ~azure.multiapi.sample.v1.models.PagingResult\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
+                    "signature": "async def _test_lro_and_paging_initial(\n    self,\n    client_request_id: Optional[str] = None,\n    test_lro_and_paging_options: Optional[_models.TestLroAndPagingOptions] = None,\n    **kwargs: Any\n) -\u003e AsyncIterator[bytes]:\n",
+                    "doc": "\"\"\"A long-running paging operation that includes a nextLink that has 10 pages.\n\n:param client_request_id: Default value is None.\n:type client_request_id: str\n:param test_lro_and_paging_options: Parameter group. Default value is None.\n:type test_lro_and_paging_options: ~azure.multiapi.sample.v1.models.TestLroAndPagingOptions\n:return: AsyncIterator[bytes] or the result of cls(response)\n:rtype: AsyncIterator[bytes]\n:raises ~azure.core.exceptions.HttpResponseError:\n\"\"\"",
                     "call": "client_request_id, test_lro_and_paging_options, **kwargs"
                 }
             },