[python] Fix import when body parameter is union of models

.chronus/changes/python-discriminated-union-2025-10-18-9-43-10.md

@@ -0,0 +1,7 @@
+---
+changeKind: fix
+packages:
+  - "@typespec/http-client-python"
+---
+
+Fix import when body parameter is union of models

packages/http-client-python/eng/scripts/ci/regenerate.ts

@@ -22,7 +22,7 @@ const argv = parseArgs({
 });
 
 // Add this near the top with other constants
-const SKIP_SPECS = ["type/union/discriminated"];
+const SKIP_SPECS: string[] = [];
 
 // Get the directory of the current file
 const PLUGIN_DIR = argv.values.pluginDir
@@ -275,6 +275,10 @@ const EMITTER_OPTIONS: Record<string, Record<string, string> | Record<string, st
     "package-name": "typetest-union",
     namespace: "typetest.union",
   },
+  "type/union/discriminated": {
+    "package-name": "typetest-discriminatedunion",
+    namespace: "typetest.discriminatedunion",
+  },
   documentation: {
     "package-name": "specs-documentation",
     namespace: "specs.documentation",

packages/http-client-python/generator/pygen/codegen/models/operation.py

@@ -408,6 +408,9 @@ def imports(  # pylint: disable=too-many-branches, disable=too-many-statements
         file_import.merge(self.get_request_builder_import(self.request_builder, async_mode, serialize_namespace))
         if self.overloads:
             file_import.add_submodule_import("typing", "overload", ImportType.STDLIB)
+            for overload in self.overloads:
+                if overload.parameters.has_body:
+                    file_import.merge(overload.parameters.body_parameter.type.imports(**kwargs))
         if self.code_model.options["models-mode"] == "dpg":
             relative_path = self.code_model.get_relative_import_path(
                 serialize_namespace, module_name="_utils.model_base"

packages/http-client-python/generator/pygen/codegen/serializers/builder_serializer.py

@@ -781,8 +781,14 @@ def _initialize_overloads(self, builder: OperationType, is_paging: bool = False)
         client_names = [
             overload.request_builder.parameters.body_parameter.client_name for overload in builder.overloads
         ]
-        for v in sorted(set(client_names), key=client_names.index):
-            retval.append(f"_{v} = None")
+        all_dpg_model_overloads = False
+        if self.code_model.options["models-mode"] == "dpg" and builder.overloads:
+            all_dpg_model_overloads = all(
+                isinstance(o.parameters.body_parameter.type, DPGModelType) for o in builder.overloads
+            )
+        if not all_dpg_model_overloads:
+            for v in sorted(set(client_names), key=client_names.index):
+                retval.append(f"_{v} = None")
         try:
             # if there is a binary overload, we do a binary check first.
             binary_overload = cast(
@@ -808,17 +814,20 @@ def _initialize_overloads(self, builder: OperationType, is_paging: bool = False)
                     f'"{other_overload.parameters.body_parameter.default_content_type}"{check_body_suffix}'
                 )
         except StopIteration:
-            for idx, overload in enumerate(builder.overloads):
-                if_statement = "if" if idx == 0 else "elif"
-                body_param = overload.parameters.body_parameter
-                retval.append(
-                    f"{if_statement} {body_param.type.instance_check_template.format(body_param.client_name)}:"
-                )
-                if body_param.default_content_type and not same_content_type:
+            if all_dpg_model_overloads:
+                retval.extend(f"{l}" for l in self._create_body_parameter(cast(OperationType, builder.overloads[0])))
+            else:
+                for idx, overload in enumerate(builder.overloads):
+                    if_statement = "if" if idx == 0 else "elif"
+                    body_param = overload.parameters.body_parameter
                     retval.append(
-                        f'    content_type = content_type or "{body_param.default_content_type}"{check_body_suffix}'
+                        f"{if_statement} {body_param.type.instance_check_template.format(body_param.client_name)}:"
                     )
-                retval.extend(f"    {l}" for l in self._create_body_parameter(cast(OperationType, overload)))
+                    if body_param.default_content_type and not same_content_type:
+                        retval.append(
+                            f'    content_type = content_type or "{body_param.default_content_type}"{check_body_suffix}'
+                        )
+                    retval.extend(f"    {l}" for l in self._create_body_parameter(cast(OperationType, overload)))
         return retval
 
     def _create_request_builder_call(

packages/http-client-python/generator/test/azure/requirements.txt

@@ -92,6 +92,7 @@ azure-mgmt-core==1.6.0
 -e ./generated/typetest-property-additionalproperties
 -e ./generated/typetest-scalar
 -e ./generated/typetest-union
+-e ./generated/typetest-discriminatedunion
 -e ./generated/typetest-model-empty
 -e ./generated/headasbooleantrue
 -e ./generated/headasbooleanfalse

packages/http-client-python/generator/test/generic_mock_api_tests/test_typetest_union_discriminated.py

@@ -0,0 +1,290 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+import pytest
+from typetest.discriminatedunion import DiscriminatedClient
+from typetest.discriminatedunion import models
+
+
+@pytest.fixture
+def client():
+    with DiscriminatedClient() as client:
+        yield client
+
+
+@pytest.fixture
+def cat_body():
+    """Cat model for testing."""
+    return models.Cat(name="Whiskers", meow=True)
+
+
+@pytest.fixture
+def dog_body():
+    """Dog model for testing."""
+    return models.Dog(name="Rex", bark=False)
+
+
+# Tests for No Envelope / Default (inline discriminator with "kind")
+@pytest.mark.skip(reason="After completely support discriminated unions, enable these tests")
+class TestNoEnvelopeDefault:
+    """Test discriminated union with inline discriminator (no envelope)."""
+
+    def test_get_default_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with default (no query param or kind=cat).
+
+        Expected response:
+        {
+            "kind": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.default.get()
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_kind_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with kind=cat query parameter.
+
+        Expected response:
+        {
+            "kind": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.default.get(kind="cat")
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_kind_dog(self, client: DiscriminatedClient, dog_body: models.Dog):
+        """Test getting dog with kind=dog query parameter.
+
+        Expected response:
+        {
+            "kind": "dog",
+            "name": "Rex",
+            "bark": false
+        }
+        """
+        result = client.no_envelope.default.get(kind="dog")
+        assert result == dog_body
+        assert isinstance(result, models.Dog)
+
+    def test_put_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test sending cat with inline discriminator.
+
+        Expected request:
+        {
+            "kind": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.default.put(cat_body)
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+
+# Tests for No Envelope / Custom Discriminator (inline with custom "type" property)
+@pytest.mark.skip(reason="After completely support discriminated unions, enable these tests")
+class TestNoEnvelopeCustomDiscriminator:
+    """Test discriminated union with inline discriminator and custom discriminator property name."""
+
+    def test_get_default_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with default (no query param or type=cat).
+
+        Expected response:
+        {
+            "type": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.custom_discriminator.get()
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_type_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with type=cat query parameter.
+
+        Expected response:
+        {
+            "type": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.custom_discriminator.get(type="cat")
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_type_dog(self, client: DiscriminatedClient, dog_body: models.Dog):
+        """Test getting dog with type=dog query parameter.
+
+        Expected response:
+        {
+            "type": "dog",
+            "name": "Rex",
+            "bark": false
+        }
+        """
+        result = client.no_envelope.custom_discriminator.get(type="dog")
+        assert result == dog_body
+        assert isinstance(result, models.Dog)
+
+    def test_put_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test sending cat with inline custom discriminator.
+
+        Expected request:
+        {
+            "type": "cat",
+            "name": "Whiskers",
+            "meow": true
+        }
+        """
+        result = client.no_envelope.custom_discriminator.put(cat_body)
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+
+# Tests for Envelope / Object / Default (envelope with "kind" and "value")
+@pytest.mark.skip(reason="After completely support discriminated unions, enable these tests")
+class TestEnvelopeObjectDefault:
+    """Test discriminated union with default envelope serialization."""
+
+    def test_get_default_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with default (no query param or kind=cat).
+
+        Expected response:
+        {
+            "kind": "cat",
+            "value": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.default.get()
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_kind_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with kind=cat query parameter.
+
+        Expected response:
+        {
+            "kind": "cat",
+            "value": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.default.get(kind="cat")
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_kind_dog(self, client: DiscriminatedClient, dog_body: models.Dog):
+        """Test getting dog with kind=dog query parameter.
+
+        Expected response:
+        {
+            "kind": "dog",
+            "value": {
+                "name": "Rex",
+                "bark": false
+            }
+        }
+        """
+        result = client.envelope.object.default.get(kind="dog")
+        assert result == dog_body
+        assert isinstance(result, models.Dog)
+
+    def test_put_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test sending cat with envelope serialization.
+
+        Expected request:
+        {
+            "kind": "cat",
+            "value": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.default.put(cat_body)
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+
+# Tests for Envelope / Object / Custom Properties (envelope with custom "petType" and "petData")
+@pytest.mark.skip(reason="After completely support discriminated unions, enable these tests")
+class TestEnvelopeObjectCustomProperties:
+    """Test discriminated union with custom property names in envelope."""
+
+    def test_get_default_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with default (no query param or petType=cat).
+
+        Expected response:
+        {
+            "petType": "cat",
+            "petData": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.custom_properties.get()
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_pet_type_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test getting cat with petType=cat query parameter.
+
+        Expected response:
+        {
+            "petType": "cat",
+            "petData": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.custom_properties.get(pet_type="cat")
+        assert result == cat_body
+        assert isinstance(result, models.Cat)
+
+    def test_get_with_pet_type_dog(self, client: DiscriminatedClient, dog_body: models.Dog):
+        """Test getting dog with petType=dog query parameter.
+
+        Expected response:
+        {
+            "petType": "dog",
+            "petData": {
+                "name": "Rex",
+                "bark": false
+            }
+        }
+        """
+        result = client.envelope.object.custom_properties.get(pet_type="dog")
+        assert result == dog_body
+        assert isinstance(result, models.Dog)
+
+    def test_put_cat(self, client: DiscriminatedClient, cat_body: models.Cat):
+        """Test sending cat with custom property names in envelope.
+
+        Expected request:
+        {
+            "petType": "cat",
+            "petData": {
+                "name": "Whiskers",
+                "meow": true
+            }
+        }
+        """
+        result = client.envelope.object.custom_properties.put(cat_body)
+        assert result == cat_body
+        assert isinstance(result, models.Cat)

packages/http-client-python/generator/test/unbranded/requirements.txt

@@ -41,6 +41,7 @@
 -e ./generated/typetest-property-additionalproperties
 -e ./generated/typetest-scalar
 -e ./generated/typetest-union
+-e ./generated/typetest-discriminatedunion
 -e ./generated/typetest-model-empty
 -e ./generated/headasbooleantrue
 -e ./generated/headasbooleanfalse