Fix deserialization error for LRO which has discriminator

.chronus/changes/deserialization-fix-2024-4-22-17-3-4.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/models/operation.py

@@ -186,6 +186,10 @@ def has_response_body(self) -> bool:
         """Tell if at least one response has a body."""
         return any(response.type for response in self.responses)
 
+    @property
+    def has_stream_kwargs(self) -> bool:
+        return self.expose_stream_keyword and self.has_response_body
+
     @property
     def any_response_has_headers(self) -> bool:
         return any(response.headers for response in self.responses)

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

@@ -562,7 +562,7 @@ def make_pipeline_call(self, builder: OperationType) -> List[str]:
         type_ignore = self.async_mode and builder.group_name == ""  # is in a mixin
         stream_value = (
             f'kwargs.pop("stream", {builder.has_stream_response})'
-            if builder.expose_stream_keyword and builder.has_response_body
+            if builder.has_stream_kwargs
             else builder.has_stream_response
         )
         return [
@@ -897,11 +897,22 @@ def _call_request_builder_helper(  # pylint: disable=too-many-statements
     def call_request_builder(self, builder: OperationType, is_paging: bool = False) -> List[str]:
         return self._call_request_builder_helper(builder, builder.request_builder, is_paging=is_paging)
 
+    @property
+    def deserialize_for_stream_res(self) -> str:
+        if self.code_model.options["version_tolerant"]:
+            return "response.iter_bytes()"
+        return (
+            "(await response.load_body()) or response._content # pylint: disable=protected-access"
+            if self.async_mode
+            else f"response.stream_download(self._client.{self.pipeline_name})"
+        )
+
     def response_headers_and_deserialization(
         self,
         builder: OperationType,
         response: Response,
     ) -> List[str]:
+        # pylint: disable=too-many-statements
         retval: List[str] = [
             (
                 f"response_headers['{response_header.wire_name}']=self._deserialize("
@@ -918,19 +929,23 @@ def response_headers_and_deserialization(
                 deserialized = f"{'await ' if self.async_mode else ''}response.read()"
             else:
                 stream_logic = False
-                if self.code_model.options["version_tolerant"]:
-                    deserialized = "response.iter_bytes()"
-                else:
-                    deserialized = f"response.stream_download(self._client.{self.pipeline_name})"
+                deserialized = self.deserialize_for_stream_res
             deserialize_code.append(f"deserialized = {deserialized}")
         elif response.type:
             pylint_disable = ""
             if isinstance(response.type, ModelType) and response.type.internal:
                 pylint_disable = "  # pylint: disable=protected-access"
             if self.code_model.options["models_mode"] == "msrest":
+                if hasattr(builder, "initial_operation") and builder.initial_operation.has_stream_kwargs:  # type: ignore # pylint: disable=line-too-long
+                    response_name = "_response"
+                    deserialize_code.append(
+                        "_response = pipeline_response if getattr(pipeline_response, 'context', {}) else pipeline_response.http_response"  # pylint: disable=line-too-long
+                    )
+                else:
+                    response_name = "pipeline_response"
                 deserialize_code.append("deserialized = self._deserialize(")
                 deserialize_code.append(f"    '{response.serialization_type}',{pylint_disable}")
-                deserialize_code.append("    pipeline_response")
+                deserialize_code.append(f"    {response_name}")
                 deserialize_code.append(")")
             elif self.code_model.options["models_mode"] == "dpg":
                 if builder.has_stream_response:
@@ -959,7 +974,7 @@ def response_headers_and_deserialization(
         if len(deserialize_code) > 0:
             if builder.expose_stream_keyword and stream_logic:
                 retval.append("if _stream:")
-                retval.append("    deserialized = response.iter_bytes()")
+                retval.append(f"    deserialized = {self.deserialize_for_stream_res}")
                 retval.append("else:")
                 retval.extend([f"    {dc}" for dc in deserialize_code])
             else:
@@ -969,11 +984,15 @@ def response_headers_and_deserialization(
     def handle_error_response(self, builder: OperationType) -> List[str]:
         async_await = "await " if self.async_mode else ""
         retval = [f"if response.status_code not in {str(builder.success_status_codes)}:"]
-        if not self.code_model.need_request_converter:
+        need_download = (
+            builder.has_stream_kwargs and self.async_mode and not self.code_model.options["version_tolerant"]
+        )
+        if not self.code_model.need_request_converter or need_download:
+            load_func = "load_body" if need_download else "read"
             retval.extend(
                 [
                     "    if _stream:",
-                    f"        {async_await} response.read()  # Load the body in memory and close the socket",
+                    f"        {async_await} response.{load_func}()  # Load the body in memory and close the socket",
                 ]
             )
         type_ignore = "  # type: ignore" if _need_type_ignore(builder) else ""
@@ -1320,6 +1339,8 @@ def initial_call(self, builder: LROOperationType) -> List[str]:
             [f"        {parameter.client_name}={parameter.client_name}," for parameter in builder.parameters.method]
         )
         retval.append("        cls=lambda x,y,z: x,")
+        if builder.initial_operation.has_stream_kwargs:
+            retval.append("        stream=True,")
         retval.append("        headers=_headers,")
         retval.append("        params=_params,")
         retval.append("        **kwargs")

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

@@ -419,6 +419,10 @@ def update_lro_operation(
             self._update_lro_operation_helper(overload)
             self.update_operation(code_model, overload["initialOperation"], is_overload=True)
 
+        # for lro initial reponse, there is no need to make deserialization so we mark it
+        # as stream operation by default which will not make deserialization by default
+        yaml_data["initialOperation"]["exposeStreamKeyword"] = True
+
     def update_paging_operation(
         self,
         code_model: Dict[str, Any],

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

@@ -109,7 +109,7 @@ def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = Non
         )
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = self._client._pipeline.run(  # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -124,10 +124,13 @@ def _basic_polling_initial(self, product: Optional[Union[JSON, IO[bytes]]] = Non
 
         deserialized = None
         if response.status_code == 200:
-            if response.content:
-                deserialized = response.json()
+            if _stream:
+                deserialized = response.iter_bytes()
             else:
-                deserialized = None
+                if response.content:
+                    deserialized = response.json()
+                else:
+                    deserialized = None
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -241,6 +244,7 @@ def begin_basic_polling(
                 product=product,
                 content_type=content_type,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs

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

@@ -82,7 +82,7 @@ async def _basic_polling_initial(
         )
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = await self._client._pipeline.run(  # type: ignore # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -97,10 +97,13 @@ async def _basic_polling_initial(
 
         deserialized = None
         if response.status_code == 200:
-            if response.content:
-                deserialized = response.json()
+            if _stream:
+                deserialized = response.iter_bytes()
             else:
-                deserialized = None
+                if response.content:
+                    deserialized = response.json()
+                else:
+                    deserialized = None
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -214,6 +217,7 @@ async def begin_basic_polling(
                 product=product,
                 content_type=content_type,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs

packages/autorest.python/samples/specification/multiapi/generated/azure/multiapi/sample/v1/aio/operations/_multiapi_service_client_operations.py

@@ -148,21 +148,26 @@ async def _test_lro_initial(
         _request = _convert_request(_request)
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = await self._client._pipeline.run(  # type: ignore # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
 
         response = pipeline_response.http_response
 
         if response.status_code not in [200, 204]:
+            if _stream:
+                await response.load_body()  # Load the body in memory and close the socket
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             error = self._deserialize.failsafe_deserialize(_models.Error, pipeline_response)
             raise HttpResponseError(response=response, model=error, error_format=ARMErrorFormat)
 
         deserialized = None
         if response.status_code == 200:
-            deserialized = self._deserialize("Product", pipeline_response)
+            if _stream:
+                deserialized = (await response.load_body()) or response._content  # pylint: disable=protected-access
+            else:
+                deserialized = self._deserialize("Product", pipeline_response)
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -230,14 +235,18 @@ async def begin_test_lro(
                 product=product,
                 content_type=content_type,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs
             )
         kwargs.pop("error_map", None)
 
         def get_long_running_output(pipeline_response):
-            deserialized = self._deserialize("Product", pipeline_response)
+            _response = (
+                pipeline_response if getattr(pipeline_response, "context", {}) else pipeline_response.http_response
+            )
+            deserialized = self._deserialize("Product", _response)
             if cls:
                 return cls(pipeline_response, deserialized, {})  # type: ignore
             return deserialized
@@ -294,18 +303,23 @@ async def _test_lro_and_paging_initial(
         _request = _convert_request(_request)
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = await self._client._pipeline.run(  # type: ignore # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
 
         response = pipeline_response.http_response
 
         if response.status_code not in [200]:
+            if _stream:
+                await response.load_body()  # Load the body in memory and close the socket
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             raise HttpResponseError(response=response, error_format=ARMErrorFormat)
 
-        deserialized = self._deserialize("PagingResult", pipeline_response)
+        if _stream:
+            deserialized = (await response.load_body()) or response._content  # pylint: disable=protected-access
+        else:
+            deserialized = self._deserialize("PagingResult", pipeline_response)
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -411,6 +425,7 @@ async def get_next(next_link=None):
                 client_request_id=client_request_id,
                 test_lro_and_paging_options=test_lro_and_paging_options,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs

packages/autorest.python/samples/specification/multiapi/generated/azure/multiapi/sample/v1/operations/_multiapi_service_client_operations.py

@@ -225,7 +225,7 @@ def _test_lro_initial(
         _request = _convert_request(_request)
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = self._client._pipeline.run(  # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -239,7 +239,10 @@ def _test_lro_initial(
 
         deserialized = None
         if response.status_code == 200:
-            deserialized = self._deserialize("Product", pipeline_response)
+            if _stream:
+                deserialized = response.stream_download(self._client._pipeline)
+            else:
+                deserialized = self._deserialize("Product", pipeline_response)
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -304,14 +307,18 @@ def begin_test_lro(
                 product=product,
                 content_type=content_type,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs
             )
         kwargs.pop("error_map", None)
 
         def get_long_running_output(pipeline_response):
-            deserialized = self._deserialize("Product", pipeline_response)
+            _response = (
+                pipeline_response if getattr(pipeline_response, "context", {}) else pipeline_response.http_response
+            )
+            deserialized = self._deserialize("Product", _response)
             if cls:
                 return cls(pipeline_response, deserialized, {})  # type: ignore
             return deserialized
@@ -368,7 +375,7 @@ def _test_lro_and_paging_initial(
         _request = _convert_request(_request)
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = self._client._pipeline.run(  # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
@@ -379,7 +386,10 @@ def _test_lro_and_paging_initial(
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             raise HttpResponseError(response=response, error_format=ARMErrorFormat)
 
-        deserialized = self._deserialize("PagingResult", pipeline_response)
+        if _stream:
+            deserialized = response.stream_download(self._client._pipeline)
+        else:
+            deserialized = self._deserialize("PagingResult", pipeline_response)
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -485,6 +495,7 @@ def get_next(next_link=None):
                 client_request_id=client_request_id,
                 test_lro_and_paging_options=test_lro_and_paging_options,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs

packages/autorest.python/test/azure/legacy/Expected/AcceptanceTests/CustomPollerPager/custompollerpager/aio/operations/_paging_operations.py

@@ -1608,18 +1608,23 @@ async def _get_multiple_pages_lro_initial(
         _request = _convert_request(_request)
         _request.url = self._client.format_url(_request.url)
 
-        _stream = False
+        _stream = kwargs.pop("stream", False)
         pipeline_response: PipelineResponse = await self._client._pipeline.run(  # pylint: disable=protected-access
             _request, stream=_stream, **kwargs
         )
 
         response = pipeline_response.http_response
 
         if response.status_code not in [202]:
+            if _stream:
+                await response.load_body()  # Load the body in memory and close the socket
             map_error(status_code=response.status_code, response=response, error_map=error_map)
             raise HttpResponseError(response=response, error_format=ARMErrorFormat)
 
-        deserialized = self._deserialize("ProductResult", pipeline_response)
+        if _stream:
+            deserialized = (await response.load_body()) or response._content  # pylint: disable=protected-access
+        else:
+            deserialized = self._deserialize("ProductResult", pipeline_response)
 
         if cls:
             return cls(pipeline_response, deserialized, {})  # type: ignore
@@ -1726,6 +1731,7 @@ async def get_next(next_link=None):
                 client_request_id=client_request_id,
                 paging_get_multiple_pages_lro_options=paging_get_multiple_pages_lro_options,
                 cls=lambda x, y, z: x,
+                stream=True,
                 headers=_headers,
                 params=_params,
                 **kwargs