Merge pull request #566 from koxudaxi/issue-480-discriminated-union-simple-types

Fix discriminated union simple types

Author: koxudaxi

docs/llms-full.txt

@@ -539,15 +539,15 @@ Run `tox run -e schema-docs` after changing supported inputs or model backends.
 
 | Format | Status | Evidence | Notes |
 |--------|--------|----------|-------|
-| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (25 fixtures) | Primary fixture format exercised under `tests/data/openapi/**/*.yaml`. |
+| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (26 fixtures) | Primary fixture format exercised under `tests/data/openapi/**/*.yaml`. |
 | OpenAPI JSON | tested | `tests/main/test_main.py::test_generate_from_json_input` | Covered by the JSON conversion CLI test in `tests/main/test_main.py`. |
 | Remote HTTP `$ref` targets | tested | `tests/main/test_main.py::test_generate_remote_ref` | Covered by the remote `$ref` generation test against a live HTTP server. |
 
 ## Fixture Suites
 
 | Suite | Fixtures | Example files | Notes |
 |-------|----------|---------------|-------|
-| Default template | 17 | `body_and_parameters.yaml`, `content_in_parameters.yaml`, `content_in_parameters_inline.yaml` | Core single-file generation scenarios exercised by the main CLI tests. |
+| Default template | 18 | `body_and_parameters.yaml`, `content_in_parameters.yaml`, `content_in_parameters_inline.yaml` | Core single-file generation scenarios exercised by the main CLI tests. |
 | Coverage fixtures | 4 | `callbacks.yaml`, `callbacks_with_operation_id.yaml`, `faux_immutability.yaml` | Focused fixtures for callbacks, non-200 responses, and other regression edges. |
 | Custom template overrides | 1 | `custom_security.yaml` | Template override coverage for `--template-dir`. |
 | Timestamp suppression | 1 | `simple.yaml` | Fixtures that exercise `--disable-timestamp`. |

docs/supported_formats.md

@@ -7,15 +7,15 @@ Run `tox run -e schema-docs` after changing supported inputs or model backends.
 
 | Format | Status | Evidence | Notes |
 |--------|--------|----------|-------|
-| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (25 fixtures) | Primary fixture format exercised under `tests/data/openapi/**/*.yaml`. |
+| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (26 fixtures) | Primary fixture format exercised under `tests/data/openapi/**/*.yaml`. |
 | OpenAPI JSON | tested | `tests/main/test_main.py::test_generate_from_json_input` | Covered by the JSON conversion CLI test in `tests/main/test_main.py`. |
 | Remote HTTP `$ref` targets | tested | `tests/main/test_main.py::test_generate_remote_ref` | Covered by the remote `$ref` generation test against a live HTTP server. |
 
 ## Fixture Suites
 
 | Suite | Fixtures | Example files | Notes |
 |-------|----------|---------------|-------|
-| Default template | 17 | `body_and_parameters.yaml`, `content_in_parameters.yaml`, `content_in_parameters_inline.yaml` | Core single-file generation scenarios exercised by the main CLI tests. |
+| Default template | 18 | `body_and_parameters.yaml`, `content_in_parameters.yaml`, `content_in_parameters_inline.yaml` | Core single-file generation scenarios exercised by the main CLI tests. |
 | Coverage fixtures | 4 | `callbacks.yaml`, `callbacks_with_operation_id.yaml`, `faux_immutability.yaml` | Focused fixtures for callbacks, non-200 responses, and other regression edges. |
 | Custom template overrides | 1 | `custom_security.yaml` | Template override coverage for `--template-dir`. |
 | Timestamp suppression | 1 | `simple.yaml` | Fixtures that exercise `--disable-timestamp`. |

fastapi_code_generator/parser.py

@@ -626,6 +626,42 @@ def parse_request_body(
         self._temporary_operation['_request'] = arguments[0] if arguments else None
         return request_body_fields
 
+    def get_field_extras(self, obj: JsonSchemaObject) -> dict[str, Any]:
+        extras = super().get_field_extras(obj)
+        if self._has_non_object_discriminator_variant(obj):
+            extras.pop('discriminator', None)
+        return extras
+
+    def _has_non_object_discriminator_variant(self, obj: JsonSchemaObject) -> bool:
+        if not obj.discriminator:
+            return False
+        combined_schemas = [*(obj.oneOf or []), *(obj.anyOf or [])]
+        if not combined_schemas:
+            return False
+        return any(
+            not self._is_object_discriminator_variant(schema)
+            for schema in combined_schemas
+        )
+
+    def _is_object_discriminator_variant(
+        self, schema: JsonSchemaObject, seen_refs: set[str] | None = None
+    ) -> bool:
+        if seen_refs is None:
+            seen_refs = set()
+        if schema.ref:
+            if schema.ref in seen_refs:
+                return False
+            seen_refs.add(schema.ref)
+            schema = JsonSchemaObject.model_validate(self.get_ref_model(schema.ref))
+        if schema.is_object or schema.properties:
+            return True
+        if not schema.allOf:
+            return False
+        return any(
+            self._is_object_discriminator_variant(member, seen_refs.copy())
+            for member in schema.allOf
+        )
+
     def _get_upload_file_type(
         self, schema: Union[JsonSchemaObject, ReferenceObject]
     ) -> tuple[str, str]:

fastapi_code_generator/prompt_data.py

@@ -509,7 +509,7 @@
     '\n'
     '| Format | Status | Evidence | Notes |\n'
     '|--------|--------|----------|-------|\n'
-    '| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (25 '
+    '| OpenAPI YAML | tested | `tests/data/openapi/**/*.yaml` (26 '
     'fixtures) | Primary fixture format exercised under '
     '`tests/data/openapi/**/*.yaml`. |\n'
     '| OpenAPI JSON | tested | '
@@ -525,7 +525,7 @@
     '\n'
     '| Suite | Fixtures | Example files | Notes |\n'
     '|-------|----------|---------------|-------|\n'
-    '| Default template | 17 | `body_and_parameters.yaml`, '
+    '| Default template | 18 | `body_and_parameters.yaml`, '
     '`content_in_parameters.yaml`, '
     '`content_in_parameters_inline.yaml` | Core single-file '
     'generation scenarios exercised by the main CLI tests. |\n'

tests/data/expected/openapi/default_template/discriminated_union_simple_type/main.py

@@ -0,0 +1,12 @@
+# generated by fastapi-codegen:
+#   filename:  discriminated_union_simple_type.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from __future__ import annotations
+
+from fastapi import FastAPI
+
+app = FastAPI(
+    title='discriminated_union_simple_type',
+    version='1.0.0',
+)

tests/data/expected/openapi/default_template/discriminated_union_simple_type/models.py

@@ -0,0 +1,35 @@
+# generated by fastapi-codegen:
+#   filename:  discriminated_union_simple_type.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Union
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class ImageUrl(BaseModel):
+    kind: str
+    image_url: str
+
+
+class AudioUrl(BaseModel):
+    kind: str
+    audio_url: str
+
+
+class DocumentUrl(BaseModel):
+    kind: str
+    document_url: str
+
+
+class BinaryContent(BaseModel):
+    kind: str
+    binary_content: str
+
+
+class Content(RootModel[Union[str, ImageUrl, AudioUrl, DocumentUrl, BinaryContent]]):
+    root: Union[str, ImageUrl, AudioUrl, DocumentUrl, BinaryContent] = Field(
+        ..., description='Union type representing different content formats'
+    )

tests/data/openapi/default_template/discriminated_union_simple_type.yaml

@@ -0,0 +1,62 @@
+info:
+  title: discriminated_union_simple_type
+  version: 1.0.0
+openapi: 3.0.0
+paths: {}
+components:
+  schemas:
+    Content:
+      oneOf:
+        - type: string
+        - $ref: '#/components/schemas/ImageUrl'
+        - $ref: '#/components/schemas/AudioUrl'
+        - $ref: '#/components/schemas/DocumentUrl'
+        - $ref: '#/components/schemas/BinaryContent'
+      discriminator:
+        propertyName: kind
+        mapping:
+          image_url: '#/components/schemas/ImageUrl'
+          audio_url: '#/components/schemas/AudioUrl'
+          document_url: '#/components/schemas/DocumentUrl'
+          binary_content: '#/components/schemas/BinaryContent'
+      description: Union type representing different content formats
+    ImageUrl:
+      type: object
+      properties:
+        kind:
+          type: string
+        image_url:
+          type: string
+      required:
+        - kind
+        - image_url
+    AudioUrl:
+      type: object
+      properties:
+        kind:
+          type: string
+        audio_url:
+          type: string
+      required:
+        - kind
+        - audio_url
+    DocumentUrl:
+      type: object
+      properties:
+        kind:
+          type: string
+        document_url:
+          type: string
+      required:
+        - kind
+        - document_url
+    BinaryContent:
+      type: object
+      properties:
+        kind:
+          type: string
+        binary_content:
+          type: string
+      required:
+        - kind
+        - binary_content

tests/main/test_main.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib.util
 import json
 from contextlib import contextmanager
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -45,6 +46,17 @@ def assert_specific_tag_routers_generated(output_dir: Path) -> None:
     validate_generated_code(output_dir)
 
 
+def assert_generated_module_has_attribute(
+    module_path: Path, module_name: str, attribute_name: str
+) -> None:
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    assert hasattr(module, attribute_name)
+
+
 def assert_generated_draft_request_is_hashable(models_path: Path) -> None:
     namespace: dict[str, Any] = {}
     exec(  # noqa: S102 - execute generated fixture code in a test namespace.
@@ -251,6 +263,24 @@ def test_generate_from_json_input(tmp_path: Path, output_dir: Path) -> None:
     validate_generated_code(output_dir)
 
 
+@freeze_time("2020-06-19")
+def test_generate_discriminated_union_with_simple_type(output_dir: Path) -> None:
+    run_cli_and_assert(
+        input_path=DATA_PATH
+        / OPEN_API_DEFAULT_TEMPLATE_DIR_NAME
+        / "discriminated_union_simple_type.yaml",
+        output_path=output_dir,
+        expected_path=EXPECTED_OPENAPI_PATH
+        / "default_template"
+        / "discriminated_union_simple_type",
+    )
+    assert_generated_module_has_attribute(
+        output_dir / "models.py",
+        "generated_discriminated_union_simple_type",
+        "Content",
+    )
+
+
 @freeze_time("2020-06-19")
 def test_generate_escapes_aliases_in_parameter_defaults(output_dir: Path) -> None:
     spec = """openapi: 3.0.0