add discriminator property support

Author: eli-bl

end_to_end_tests/__init__.py

@@ -1 +1,4 @@
 """ Generate a complete client and verify that it is correct """
+import pytest
+
+pytest.register_assert_rewrite('end_to_end_tests.end_to_end_live_tests')

end_to_end_tests/baseline_openapi_3.0.json

@@ -2841,15 +2841,17 @@
           "modelType": {
             "type": "string"
           }
-        }
+        },
+        "required": ["modelType"]
       },
       "ADiscriminatedUnionType2": {
         "type": "object",
         "properties": {
           "modelType": {
             "type": "string"
           }
-        }
+        },
+        "required": ["modelType"]
       }
     },
     "parameters": {

end_to_end_tests/baseline_openapi_3.1.yaml

@@ -2835,15 +2835,17 @@ info:
         "modelType": {
           "type": "string"
         }
-      }
+      },
+      "required": ["modelType"]
     },
     "ADiscriminatedUnionType2": {
       "type": "object",
       "properties": {
         "modelType": {
           "type": "string"
         }
-      }
+      },
+      "required": ["modelType"]
     }
   }
   "parameters": {

end_to_end_tests/end_to_end_live_tests.py

@@ -0,0 +1,33 @@
+import importlib
+from typing import Any
+
+import pytest
+
+
+def live_tests_3_x():
+    _test_model_with_discriminated_union()
+
+
+def _import_model(module_name, class_name: str) -> Any:
+    module = importlib.import_module(f"my_test_api_client.models.{module_name}")
+    module = importlib.reload(module)  # avoid test contamination from previous import
+    return getattr(module, class_name)
+
+
+def _test_model_with_discriminated_union():
+    ModelType1Class = _import_model("a_discriminated_union_type_1", "ADiscriminatedUnionType1")
+    ModelType2Class = _import_model("a_discriminated_union_type_2", "ADiscriminatedUnionType2")
+    ModelClass = _import_model("model_with_discriminated_union", "ModelWithDiscriminatedUnion")
+
+    assert (
+        ModelClass.from_dict({"discriminated_union": {"modelType": "type1"}}) ==
+        ModelClass(discriminated_union=ModelType1Class.from_dict({"modelType": "type1"}))
+    )
+    assert (
+        ModelClass.from_dict({"discriminated_union": {"modelType": "type2"}}) ==
+        ModelClass(discriminated_union=ModelType2Class.from_dict({"modelType": "type2"}))
+    )
+    with pytest.raises(TypeError):
+        ModelClass.from_dict({"discriminated_union": {"modelType": "type3"}})
+    with pytest.raises(TypeError):
+        ModelClass.from_dict({"discriminated_union": {}})

end_to_end_tests/golden-record/my_test_api_client/models/a_discriminated_union_type_1.py

@@ -1,38 +1,38 @@
-from typing import Any, Dict, List, Type, TypeVar, Union
+from typing import Any, Dict, List, Type, TypeVar
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
-from ..types import UNSET, Unset
-
 T = TypeVar("T", bound="ADiscriminatedUnionType1")
 
 
 @_attrs_define
 class ADiscriminatedUnionType1:
     """
     Attributes:
-        model_type (Union[Unset, str]):
+        model_type (str):
     """
 
-    model_type: Union[Unset, str] = UNSET
+    model_type: str
     additional_properties: Dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> Dict[str, Any]:
         model_type = self.model_type
 
         field_dict: Dict[str, Any] = {}
         field_dict.update(self.additional_properties)
-        field_dict.update({})
-        if model_type is not UNSET:
-            field_dict["modelType"] = model_type
+        field_dict.update(
+            {
+                "modelType": model_type,
+            }
+        )
 
         return field_dict
 
     @classmethod
     def from_dict(cls: Type[T], src_dict: Dict[str, Any]) -> T:
         d = src_dict.copy()
-        model_type = d.pop("modelType", UNSET)
+        model_type = d.pop("modelType")
 
         a_discriminated_union_type_1 = cls(
             model_type=model_type,

end_to_end_tests/golden-record/my_test_api_client/models/a_discriminated_union_type_2.py

@@ -1,38 +1,38 @@
-from typing import Any, Dict, List, Type, TypeVar, Union
+from typing import Any, Dict, List, Type, TypeVar
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
-from ..types import UNSET, Unset
-
 T = TypeVar("T", bound="ADiscriminatedUnionType2")
 
 
 @_attrs_define
 class ADiscriminatedUnionType2:
     """
     Attributes:
-        model_type (Union[Unset, str]):
+        model_type (str):
     """
 
-    model_type: Union[Unset, str] = UNSET
+    model_type: str
     additional_properties: Dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> Dict[str, Any]:
         model_type = self.model_type
 
         field_dict: Dict[str, Any] = {}
         field_dict.update(self.additional_properties)
-        field_dict.update({})
-        if model_type is not UNSET:
-            field_dict["modelType"] = model_type
+        field_dict.update(
+            {
+                "modelType": model_type,
+            }
+        )
 
         return field_dict
 
     @classmethod
     def from_dict(cls: Type[T], src_dict: Dict[str, Any]) -> T:
         d = src_dict.copy()
-        model_type = d.pop("modelType", UNSET)
+        model_type = d.pop("modelType")
 
         a_discriminated_union_type_2 = cls(
             model_type=model_type,

end_to_end_tests/golden-record/my_test_api_client/models/model_with_discriminated_union.py

@@ -59,23 +59,34 @@ def _parse_discriminated_union(
                 return data
             if isinstance(data, Unset):
                 return data
-            try:
-                if not isinstance(data, dict):
-                    raise TypeError()
-                componentsschemas_a_discriminated_union_type_0 = ADiscriminatedUnionType1.from_dict(data)
-
-                return componentsschemas_a_discriminated_union_type_0
-            except:  # noqa: E722
-                pass
-            try:
-                if not isinstance(data, dict):
-                    raise TypeError()
-                componentsschemas_a_discriminated_union_type_1 = ADiscriminatedUnionType2.from_dict(data)
-
-                return componentsschemas_a_discriminated_union_type_1
-            except:  # noqa: E722
-                pass
-            return cast(Union["ADiscriminatedUnionType1", "ADiscriminatedUnionType2", None, Unset], data)
+            if not isinstance(data, dict):
+                raise TypeError()
+            if "modelType" in data:
+                _discriminator_value = data["modelType"]
+
+                def _parse_componentsschemas_a_discriminated_union_type_1(data: object) -> ADiscriminatedUnionType1:
+                    if not isinstance(data, dict):
+                        raise TypeError()
+                    componentsschemas_a_discriminated_union_type_1 = ADiscriminatedUnionType1.from_dict(data)
+
+                    return componentsschemas_a_discriminated_union_type_1
+
+                def _parse_componentsschemas_a_discriminated_union_type_2(data: object) -> ADiscriminatedUnionType2:
+                    if not isinstance(data, dict):
+                        raise TypeError()
+                    componentsschemas_a_discriminated_union_type_2 = ADiscriminatedUnionType2.from_dict(data)
+
+                    return componentsschemas_a_discriminated_union_type_2
+
+                _discriminator_mapping = {
+                    "type1": _parse_componentsschemas_a_discriminated_union_type_1,
+                    "type2": _parse_componentsschemas_a_discriminated_union_type_2,
+                }
+                if _parse_fn := _discriminator_mapping.get(_discriminator_value):
+                    return cast(
+                        Union["ADiscriminatedUnionType1", "ADiscriminatedUnionType2", None, Unset], _parse_fn(data)
+                    )
+            raise TypeError("unrecognized value for property modelType")
 
         discriminated_union = _parse_discriminated_union(d.pop("discriminated_union", UNSET))
 

end_to_end_tests/test_end_to_end.py

@@ -1,13 +1,17 @@
+import os
 import shutil
 from filecmp import cmpfiles, dircmp
 from pathlib import Path
-from typing import Dict, List, Optional, Set
+import sys
+from typing import Callable, Dict, List, Optional, Set
 
 import pytest
 from click.testing import Result
 from typer.testing import CliRunner
 
 from openapi_python_client.cli import app
+from .end_to_end_live_tests import live_tests_3_x
+
 
 
 def _compare_directories(
@@ -83,8 +87,10 @@ def run_e2e_test(
     golden_record_path: str = "golden-record",
     output_path: str = "my-test-api-client",
     expected_missing: Optional[Set[str]] = None,
+    live_tests: Optional[Callable[[str], None]] = None,
 ) -> Result:
-    output_path = Path.cwd() / output_path
+    cwd = Path.cwd()
+    output_path = cwd / output_path
     shutil.rmtree(output_path, ignore_errors=True)
     result = generate(extra_args, openapi_document)
     gr_path = Path(__file__).parent / golden_record_path
@@ -97,6 +103,13 @@ def run_e2e_test(
     _compare_directories(
         gr_path, output_path, expected_differences=expected_differences, expected_missing=expected_missing
     )
+    if live_tests:
+        old_path = sys.path.copy()
+        sys.path.insert(0, str(output_path))
+        try:
+            live_tests()
+        finally:
+            sys.path = old_path
 
     import mypy.api
 
@@ -131,11 +144,11 @@ def _run_command(command: str, extra_args: Optional[List[str]] = None, openapi_d
 
 
 def test_baseline_end_to_end_3_0():
-    run_e2e_test("baseline_openapi_3.0.json", [], {})
+    run_e2e_test("baseline_openapi_3.0.json", [], {}, live_tests=live_tests_3_x)
 
 
 def test_baseline_end_to_end_3_1():
-    run_e2e_test("baseline_openapi_3.1.yaml", [], {})
+    run_e2e_test("baseline_openapi_3.1.yaml", [], {}, live_tests=live_tests_3_x)
 
 
 def test_3_1_specific_features():

openapi_python_client/parser/properties/union.py

@@ -10,7 +10,15 @@
 from ...utils import PythonIdentifier
 from ..errors import ParseError, PropertyError
 from .protocol import PropertyProtocol, Value
-from .schemas import Schemas
+from .schemas import Schemas, parse_reference_path
+
+
+@define
+class DiscriminatorDefinition:
+    from .model_property import ModelProperty
+
+    property_name: str
+    value_to_model_map: dict[str, ModelProperty]
 
 
 @define
@@ -24,6 +32,7 @@ class UnionProperty(PropertyProtocol):
     description: str | None
     example: str | None
     inner_properties: list[PropertyProtocol]
+    discriminators: list[DiscriminatorDefinition] | None = None
     template: ClassVar[str] = "union_property.py.jinja"
 
     @classmethod
@@ -67,16 +76,7 @@ def build(
                 return PropertyError(detail=f"Invalid property in union {name}", data=sub_prop_data), schemas
             sub_properties.append(sub_prop)
 
-        def flatten_union_properties(sub_properties: list[PropertyProtocol]) -> list[PropertyProtocol]:
-            flattened = []
-            for sub_prop in sub_properties:
-                if isinstance(sub_prop, UnionProperty):
-                    flattened.extend(flatten_union_properties(sub_prop.inner_properties))
-                else:
-                    flattened.append(sub_prop)
-            return flattened
-
-        sub_properties = flatten_union_properties(sub_properties)
+        sub_properties, discriminators_list = _flatten_union_properties(sub_properties)
 
         prop = UnionProperty(
             name=name,
@@ -92,6 +92,17 @@ def flatten_union_properties(sub_properties: list[PropertyProtocol]) -> list[Pro
             default_or_error.data = data
             return default_or_error, schemas
         prop = evolve(prop, default=default_or_error)
+
+        if data.discriminator:
+            discriminator_or_error = _parse_discriminator(data.discriminator, sub_properties, schemas)
+            if isinstance(discriminator_or_error, PropertyError):
+                return discriminator_or_error, schemas
+            discriminators_list = [discriminator_or_error, *discriminators_list]
+        if discriminators_list:
+            if error := _validate_discriminators(discriminators_list):
+                return error, schemas
+            prop = evolve(prop, discriminators=discriminators_list)
+
         return prop, schemas
 
     def convert_value(self, value: Any) -> Value | None | PropertyError:
@@ -189,3 +200,103 @@ def validate_location(self, location: oai.ParameterLocation) -> ParseError | Non
             if evolve(cast(Property, inner_prop), required=self.required).validate_location(location) is not None:
                 return ParseError(detail=f"{self.get_type_string()} is not allowed in {location}")
         return None
+
+
+def _flatten_union_properties(
+    sub_properties: list[PropertyProtocol],
+) -> tuple[list[PropertyProtocol], list[DiscriminatorDefinition]]:
+    flattened = []
+    discriminators = []
+    for sub_prop in sub_properties:
+        if isinstance(sub_prop, UnionProperty):
+            if sub_prop.discriminators:
+                discriminators.extend(sub_prop.discriminators)
+            new_props, new_discriminators = _flatten_union_properties(sub_prop.inner_properties)
+            flattened.extend(new_props)
+            discriminators.extend(new_discriminators)
+        else:
+            flattened.append(sub_prop)
+    return flattened, discriminators
+
+
+def _parse_discriminator(
+    data: oai.Discriminator,
+    subtypes: list[PropertyProtocol],
+    schemas: Schemas,
+) -> DiscriminatorDefinition | PropertyError:
+    from .model_property import ModelProperty
+
+    # See: https://spec.openapis.org/oas/v3.1.0.html#discriminator-object
+
+    def _find_top_level_model(matching_model: ModelProperty) -> ModelProperty | None:
+        # This is needed because, when we built the union list, $refs were changed into a copy of
+        # the type they referred to, without preserving the original name. We need to know that
+        # every type in the discriminator is a $ref to a top-level type and we need its name.
+        for prop in schemas.classes_by_reference.values():
+            if isinstance(prop, ModelProperty):
+                if prop.class_info == matching_model.class_info:
+                    return prop
+        return None
+
+    model_types_by_name: dict[str, ModelProperty] = {}
+    for model in subtypes:
+        # Note, model here can never be a UnionProperty, because we've already done
+        # flatten_union_properties() before this point.
+        if not isinstance(model, ModelProperty):
+            return PropertyError(
+                detail="All schema variants must be objects when using a discriminator",
+            )
+        top_level_model = _find_top_level_model(model)
+        if not top_level_model:
+            return PropertyError(
+                detail="Inline schema declarations are not allowed when using a discriminator",
+            )
+        name = top_level_model.name
+        if name.startswith("/components/schemas/"):
+            name = name.split("/", 3)[-1]
+        model_types_by_name[name] = top_level_model
+
+    # The discriminator can specify an explicit mapping of values to types, but it doesn't
+    # have to; the default behavior is that the value for each type is simply its name.
+    mapping: dict[str, ModelProperty] = model_types_by_name.copy()
+    if data.mapping:
+        for discriminator_value, model_ref in data.mapping.items():
+            ref_path = parse_reference_path(
+                model_ref if model_ref.startswith("#/components/schemas/") else f"#/components/schemas/{model_ref}"
+            )
+            if isinstance(ref_path, ParseError) or ref_path not in schemas.classes_by_reference:
+                return PropertyError(detail=f'Invalid reference "{model_ref}" in discriminator mapping')
+            name = ref_path.split("/", 3)[-1]
+            if not (lookup_model := model_types_by_name.get(name)):
+                return PropertyError(
+                    detail=f'Discriminator mapping referred to "{model_ref}" which is not one of the schema variants',
+                )
+            for original_value in (name for name, m in model_types_by_name.items() if m == lookup_model):
+                mapping.pop(original_value)
+            mapping[discriminator_value] = lookup_model
+    else:
+        mapping = model_types_by_name
+
+    return DiscriminatorDefinition(property_name=data.propertyName, value_to_model_map=mapping)
+
+
+def _validate_discriminators(
+    discriminators: list[DiscriminatorDefinition],
+) -> PropertyError | None:
+    prop_names_values_classes = [
+        (discriminator.property_name, key, model.class_info.name)
+        for discriminator in discriminators
+        for key, model in discriminator.value_to_model_map.items()
+    ]
+    for p, v in {(p, v) for p, v, _ in prop_names_values_classes}:
+        if len({c for p1, v1, c in prop_names_values_classes if (p1, v1) == (p, v)}) > 1:
+            return PropertyError(f'Discriminator property "{p}" had more than one schema for value "{v}"')
+    return None
+
+    # TODO: We should also validate that property_name refers to a property that 1. exists,
+    # 2. is required, 3. is a string (in all of these models). However, currently we can't
+    # do that because, at the time this function is called, the ModelProperties within the
+    # union haven't yet been post-processed and so we don't have full information about
+    # their properties. To fix this, we may need to generalize the post-processing phase so
+    # that any Property type, not just ModelProperty, can say it needs post-processing; then
+    # we can defer _validate_discriminators till that phase.

openapi_python_client/schema/__init__.py

@@ -1,4 +1,5 @@
 __all__ = [
+    "Discriminator",
     "MediaType",
     "OpenAPI",
     "Operation",
@@ -17,6 +18,7 @@
 
 from .data_type import DataType
 from .openapi_schema_pydantic import (
+    Discriminator,
     MediaType,
     OpenAPI,
     Operation,

openapi_python_client/templates/property_templates/union_property.py.jinja

@@ -1,3 +1,36 @@
+{% macro construct_inner_property(inner_property) %}
+{% import "property_templates/" + inner_property.template as inner_template %}
+{% if inner_template.check_type_for_construct %}
+if not {{ inner_template.check_type_for_construct(inner_property, "data") }}:
+    raise TypeError()
+{% endif %}
+{{ inner_template.construct(inner_property, "data") }}
+return {{ inner_property.python_name }}
+{%- endmacro %}
+
+{% macro construct_discriminator_lookup(property) %}
+{% set _discriminator_properties = [] -%}
+{% for discriminator in property.discriminators %}
+{{- _discriminator_properties.append(discriminator.property_name) or "" -}}
+if not isinstance(data, dict):
+    raise TypeError()
+if "{{ discriminator.property_name }}" in data:
+    _discriminator_value = data["{{ discriminator.property_name }}"]
+    {% for model in discriminator.value_to_model_map.values() %}
+    def _parse_{{ model.python_name }}(data: object) -> {{ model.get_type_string() }}:
+{{ construct_inner_property(model) | indent(12, True) }}
+    {% endfor %}
+    _discriminator_mapping = {
+    {% for value, model in discriminator.value_to_model_map.items() %}
+        "{{ value }}": _parse_{{ model.python_name }},
+    {% endfor %}
+    }
+    if _parse_fn := _discriminator_mapping.get(_discriminator_value):
+        return cast({{ property.get_type_string() }}, _parse_fn(data))
+{% endfor %}
+raise TypeError(f"unrecognized value for property {{ _discriminator_properties | join(' or ') }}")
+{% endmacro %}
+
 {% macro construct(property, source) %}
 def _parse_{{ property.python_name }}(data: object) -> {{ property.get_type_string() }}:
     {% if "None" in property.get_type_strings_in_union(json=True, multipart=False) %}
@@ -8,6 +41,9 @@ def _parse_{{ property.python_name }}(data: object) -> {{ property.get_type_stri
     if isinstance(data, Unset):
         return data
     {% endif %}
+{% if property.discriminators %}
+{{ construct_discriminator_lookup(property) | indent(4, True) }}
+{% else %}
     {% set ns = namespace(contains_unmodified_properties = false) %}
     {% for inner_property in property.inner_properties %}
     {% import "property_templates/" + inner_property.template as inner_template %}
@@ -17,24 +53,17 @@ def _parse_{{ property.python_name }}(data: object) -> {{ property.get_type_stri
         {% endif %}
     {% if inner_template.check_type_for_construct and (not loop.last or ns.contains_unmodified_properties) %}
     try:
-        if not {{ inner_template.check_type_for_construct(inner_property, "data") }}:
-            raise TypeError()
-        {{ inner_template.construct(inner_property, "data") | indent(8) }}
-        return {{ inner_property.python_name }}
+{{ construct_inner_property(inner_property) | indent(8, True) }}
     except: # noqa: E722
         pass
     {% else  %}{# Don't do try/except for the last one nor any properties with no type checking #}
-    {% if inner_template.check_type_for_construct %}
-    if not {{ inner_template.check_type_for_construct(inner_property, "data") }}:
-        raise TypeError()
-    {% endif %}
-    {{ inner_template.construct(inner_property, "data") | indent(4) }}
-    return {{ inner_property.python_name }}
+{{ construct_inner_property(inner_property) | indent(4, True) }}
     {% endif %}
     {% endfor %}
     {% if ns.contains_unmodified_properties %}
     return cast({{ property.get_type_string() }}, data)
     {% endif %}
+{% endif %}
 
 {{ property.python_name }} = _parse_{{ property.python_name }}({{ source }})
 {% endmacro %}

tests/test_parser/test_properties/properties_test_helpers.py

@@ -0,0 +1,16 @@
+import re
+from typing import Any, Union
+
+from openapi_python_client.parser.errors import PropertyError
+from openapi_python_client.parser.properties.property import Property
+
+
+def assert_prop_error(
+    p: Union[Property, PropertyError],
+    message_regex: str,
+    data: Any = None,
+) -> None:
+    assert isinstance(p, PropertyError)
+    assert re.search(message_regex, p.detail)
+    if data is not None:
+        assert p.data == data

tests/test_parser/test_properties/test_init.py

@@ -631,7 +631,7 @@ def test_property_from_data_ref_model(self, model_property_factory, config):
             required=required,
             class_info=class_info,
         )
-        assert schemas == new_schemas
+        assert new_schemas == schemas
 
     def test_property_from_data_ref_not_found(self, mocker):
         from openapi_python_client.parser.properties import PropertyError, Schemas, property_from_data

tests/test_parser/test_properties/test_union.py

@@ -1,8 +1,18 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import pytest
+from attr import evolve
+
 import openapi_python_client.schema as oai
+from openapi_python_client.config import Config
 from openapi_python_client.parser.errors import ParseError, PropertyError
 from openapi_python_client.parser.properties import Schemas, UnionProperty
+from openapi_python_client.parser.properties.model_property import ModelProperty
+from openapi_python_client.parser.properties.property import Property
 from openapi_python_client.parser.properties.protocol import Value
+from openapi_python_client.parser.properties.schemas import Class
 from openapi_python_client.schema import DataType, ParameterLocation
+from tests.test_parser.test_properties.properties_test_helpers import assert_prop_error
 
 
 def test_property_from_data_union(union_property_factory, date_time_property_factory, string_property_factory, config):
@@ -33,6 +43,206 @@ def test_property_from_data_union(union_property_factory, date_time_property_fac
     assert s == Schemas()
 
 
+def _make_basic_model(
+    name: str,
+    props: Dict[str, oai.Schema],
+    required_prop: Optional[str],
+    schemas: Schemas,
+    config: Config,
+) -> Tuple[ModelProperty, Schemas]:
+    model, schemas = ModelProperty.build(
+        data=oai.Schema.model_construct(
+            required=[required_prop] if required_prop else [],
+            title=name,
+            properties=props,
+        ),
+        name=name or "some_generated_name",
+        schemas=schemas,
+        required=False,
+        parent_name="",
+        config=config,
+        roots={"root"},
+        process_properties=True,
+    )
+    assert isinstance(model, ModelProperty)
+    if name:
+        schemas = evolve(
+            schemas, classes_by_reference={**schemas.classes_by_reference, f"/components/schemas/{name}": model}
+        )
+    return model, schemas
+
+
+def _assert_valid_discriminator(
+    p: Union[Property, PropertyError],
+    expected_discriminators: List[Tuple[str, Dict[str, Class]]],
+) -> None:
+    assert isinstance(p, UnionProperty)
+    assert p.discriminators
+    assert [(d[0], {key: model.class_info for key, model in d[1].items()}) for d in expected_discriminators] == [
+        (d.property_name, {key: model.class_info for key, model in d.value_to_model_map.items()})
+        for d in p.discriminators
+    ]
+
+
+def test_discriminator_with_explicit_mapping(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Reference(ref="#/components/schemas/Model2"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+            mapping={
+                # mappings can use either a fully-qualified schema reference or just the schema name
+                "type1": "#/components/schemas/Model1",
+                "type2": "Model2",
+            },
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    _assert_valid_discriminator(p, [("type", {"type1": model1, "type2": model2})])
+
+
+def test_discriminator_with_implicit_mapping(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Reference(ref="#/components/schemas/Model2"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    _assert_valid_discriminator(p, [("type", {"Model1": model1, "Model2": model2})])
+
+
+def test_discriminator_with_partial_explicit_mapping(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Reference(ref="#/components/schemas/Model2"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+            mapping={
+                "type1": "#/components/schemas/Model1",
+                # no value specified for Model2, so it defaults to just "Model2"
+            },
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    _assert_valid_discriminator(p, [("type", {"type1": model1, "Model2": model2})])
+
+
+def test_discriminators_in_nested_unions_same_property(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    model3, schemas = _make_basic_model("Model3", props, "type", schemas, config)
+    model4, schemas = _make_basic_model("Model4", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model1"),
+                    oai.Reference(ref="#/components/schemas/Model2"),
+                ],
+                discriminator=oai.Discriminator.model_construct(propertyName="type"),
+            ),
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model3"),
+                    oai.Reference(ref="#/components/schemas/Model4"),
+                ],
+                discriminator=oai.Discriminator.model_construct(propertyName="type"),
+            ),
+        ],
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    _assert_valid_discriminator(
+        p,
+        [
+            ("type", {"Model1": model1, "Model2": model2}),
+            ("type", {"Model3": model3, "Model4": model4}),
+        ],
+    )
+
+
+def test_discriminators_in_nested_unions_different_property(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props1 = {"type": oai.Schema.model_construct(type="string")}
+    props2 = {"other": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props1, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props1, "type", schemas, config)
+    model3, schemas = _make_basic_model("Model3", props2, "other", schemas, config)
+    model4, schemas = _make_basic_model("Model4", props2, "other", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model1"),
+                    oai.Reference(ref="#/components/schemas/Model2"),
+                ],
+                discriminator=oai.Discriminator.model_construct(propertyName="type"),
+            ),
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model3"),
+                    oai.Reference(ref="#/components/schemas/Model4"),
+                ],
+                discriminator=oai.Discriminator.model_construct(propertyName="other"),
+            ),
+        ],
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    _assert_valid_discriminator(
+        p,
+        [
+            ("type", {"Model1": model1, "Model2": model2}),
+            ("other", {"Model3": model3, "Model4": model4}),
+        ],
+    )
+
+
 def test_build_union_property_invalid_property(config):
     name = "bad_union"
     required = True
@@ -42,7 +252,7 @@ def test_build_union_property_invalid_property(config):
     p, s = UnionProperty.build(
         name=name, required=required, data=data, schemas=Schemas(), parent_name="parent", config=config
     )
-    assert p == PropertyError(detail=f"Invalid property in union {name}", data=reference)
+    assert_prop_error(p, f"Invalid property in union {name}", data=reference)
 
 
 def test_invalid_default(config):
@@ -82,3 +292,155 @@ def test_not_required_in_path(config):
 
     err = prop.validate_location(ParameterLocation.PATH)
     assert isinstance(err, ParseError)
+
+
+@pytest.mark.parametrize("bad_ref", ["#/components/schemas/UnknownModel", "http://remote/Model2"])
+def test_discriminator_invalid_reference(bad_ref, config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Reference(ref="#/components/schemas/Model2"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+            mapping={
+                "Model1": "#/components/schemas/Model1",
+                "Model2": bad_ref,
+            },
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    assert_prop_error(p, "^Invalid reference")
+
+
+def test_discriminator_mapping_uses_schema_not_in_list(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    model3, schemas = _make_basic_model("Model3", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Reference(ref="#/components/schemas/Model2"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+            mapping={
+                "Model1": "#/components/schemas/Model1",
+                "Model3": "#/components/schemas/Model3",
+            },
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    assert_prop_error(p, "not one of the schema variants")
+
+
+def test_discriminator_invalid_variant_is_not_object(config, string_property_factory):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model_type, schemas = _make_basic_model("ModelType", props, "type", schemas, config)
+    string_type = string_property_factory()
+    schemas = evolve(
+        schemas,
+        classes_by_reference={
+            **schemas.classes_by_reference,
+            "/components/schemas/StringType": string_type,
+        },
+    )
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/ModelType"),
+            oai.Reference(ref="#/components/schemas/StringType"),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    assert_prop_error(p, "must be objects")
+
+
+def test_discriminator_invalid_inline_schema_variant(config, string_property_factory):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Reference(ref="#/components/schemas/Model1"),
+            oai.Schema.model_construct(
+                type="object",
+                properties=props,
+            ),
+        ],
+        discriminator=oai.Discriminator.model_construct(
+            propertyName="type",
+        ),
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    assert_prop_error(p, "Inline schema")
+
+
+def test_conflicting_discriminator_mappings(config):
+    from openapi_python_client.parser.properties import Schemas, property_from_data
+
+    schemas = Schemas()
+    props = {"type": oai.Schema.model_construct(type="string")}
+    model1, schemas = _make_basic_model("Model1", props, "type", schemas, config)
+    model2, schemas = _make_basic_model("Model2", props, "type", schemas, config)
+    model3, schemas = _make_basic_model("Model3", props, "type", schemas, config)
+    model4, schemas = _make_basic_model("Model4", props, "type", schemas, config)
+    data = oai.Schema.model_construct(
+        oneOf=[
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model1"),
+                    oai.Reference(ref="#/components/schemas/Model2"),
+                ],
+                discriminator=oai.Discriminator.model_construct(
+                    propertyName="type",
+                    mapping={"a": "Model1", "b": "Model2"},
+                ),
+            ),
+            oai.Schema.model_construct(
+                oneOf=[
+                    oai.Reference(ref="#/components/schemas/Model3"),
+                    oai.Reference(ref="#/components/schemas/Model4"),
+                ],
+                discriminator=oai.Discriminator.model_construct(
+                    propertyName="type",
+                    mapping={"a": "Model3", "x": "Model4"},
+                ),
+            ),
+        ],
+    )
+
+    p, schemas = property_from_data(
+        name="MyUnion", required=False, data=data, schemas=schemas, parent_name="parent", config=config
+    )
+    assert_prop_error(p, '"type" had more than one schema for value "a"')