simply the discriminator logic a bit

Author: eli-bl

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

@@ -64,23 +64,23 @@ def _parse_discriminated_union(
             if "modelType" in data:
                 _discriminator_value = data["modelType"]
 
-                def _parse_componentsschemas_a_discriminated_union_type_1(data: object) -> ADiscriminatedUnionType1:
+                def _parse_1(data: object) -> ADiscriminatedUnionType1:
                     if not isinstance(data, dict):
                         raise TypeError()
-                    componentsschemas_a_discriminated_union_type_1 = ADiscriminatedUnionType1.from_dict(data)
+                    componentsschemas_a_discriminated_union_type_0 = ADiscriminatedUnionType1.from_dict(data)
 
-                    return componentsschemas_a_discriminated_union_type_1
+                    return componentsschemas_a_discriminated_union_type_0
 
-                def _parse_componentsschemas_a_discriminated_union_type_2(data: object) -> ADiscriminatedUnionType2:
+                def _parse_2(data: object) -> ADiscriminatedUnionType2:
                     if not isinstance(data, dict):
                         raise TypeError()
-                    componentsschemas_a_discriminated_union_type_2 = ADiscriminatedUnionType2.from_dict(data)
+                    componentsschemas_a_discriminated_union_type_1 = ADiscriminatedUnionType2.from_dict(data)
 
-                    return componentsschemas_a_discriminated_union_type_2
+                    return componentsschemas_a_discriminated_union_type_1
 
                 _discriminator_mapping = {
-                    "type1": _parse_componentsschemas_a_discriminated_union_type_1,
-                    "type2": _parse_componentsschemas_a_discriminated_union_type_2,
+                    "type1": _parse_1,
+                    "type2": _parse_2,
                 }
                 if _parse_fn := _discriminator_mapping.get(_discriminator_value):
                     return cast(

openapi_python_client/parser/properties/model_property.py

@@ -32,6 +32,7 @@ class ModelProperty(PropertyProtocol):
     relative_imports: set[str] | None
     lazy_imports: set[str] | None
     additional_properties: Property | None
+    ref_path: ReferencePath | None = None
     _json_type_string: ClassVar[str] = "Dict[str, Any]"
 
     template: ClassVar[str] = "model_property.py.jinja"

openapi_python_client/parser/properties/protocol.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from openapi_python_client.parser.properties.schemas import ReferencePath
+
 __all__ = ["PropertyProtocol", "Value"]
 
 from abc import abstractmethod
@@ -185,3 +187,6 @@ def is_base_type(self) -> bool:
             ListProperty.__name__,
             UnionProperty.__name__,
         }
+
+    def get_ref_path(self) -> ReferencePath | None:
+        return self.ref_path if hasattr(self, "ref_path") else None

openapi_python_client/parser/properties/schemas.py

@@ -142,6 +142,15 @@ def update_schemas_with_data(
             )
         return prop
 
+    # Save the original path (/components/schemas/X) in the property. This is important because:
+    # 1. There are some contexts (such as a union with a discriminator) where we have a Property
+    #    instance and we want to know what its path is, instead of the other way round.
+    # 2. Even though we did set prop.name to be the same as ref_path when we created it above,
+    #    whenever there's a $ref to this property, we end up making a copy of it and changing
+    #    the name. So we can't rely on prop.name always being the path.
+    if hasattr(prop, "ref_path"):
+        prop.ref_path = ref_path
+
     schemas = evolve(schemas, classes_by_reference={ref_path: prop, **schemas.classes_by_reference})
     return schemas
 

openapi_python_client/parser/properties/union.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from itertools import chain
-from typing import Any, ClassVar, cast
+from typing import Any, ClassVar, OrderedDict, cast
 
 from attr import define, evolve
 
@@ -15,6 +15,27 @@
 
 @define
 class DiscriminatorDefinition:
+    """Represents a discriminator that can optionally be specified for a union type.
+    
+    Normally, a UnionProperty has either zero or one of these. However, a nested union
+    could have more than one, as we accumulate all the discriminators when we flatten
+    out the nested schemas. For example:
+
+        anyOf:
+          - anyOf:
+              - $ref: "#/components/schemas/Cat"
+              - $ref: "#/components/schemas/Dog"
+            discriminator:
+              propertyName: mammalType
+          - anyOf:
+              - $ref: "#/components/schemas/Condor"
+              - $ref: "#/components/schemas/Chicken"
+            discriminator:
+              propertyName: birdType
+
+    In this example there are four schemas and two discriminators. The deserializer
+    logic will check for the mammalType property first, then birdType.
+    """
     property_name: str
     value_to_model_map: dict[str, PropertyProtocol]
     # Every value in the map is really a ModelProperty, but this avoids circular imports
@@ -75,7 +96,7 @@ def build(
                 return PropertyError(detail=f"Invalid property in union {name}", data=sub_prop_data), schemas
             sub_properties.append(sub_prop)
 
-        sub_properties, discriminators_list = _flatten_union_properties(sub_properties)
+        sub_properties, discriminators_from_nested_unions = _flatten_union_properties(sub_properties)
 
         prop = UnionProperty(
             name=name,
@@ -92,15 +113,14 @@ def build(
             return default_or_error, schemas
         prop = evolve(prop, default=default_or_error)
 
+        all_discriminators = discriminators_from_nested_unions
         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)
+            all_discriminators = [discriminator_or_error, *all_discriminators]
+        if all_discriminators:
+            prop = evolve(prop, discriminators=all_discriminators)
 
         return prop, schemas
 
@@ -227,15 +247,33 @@ def _parse_discriminator(
 
     # 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
+    # Conditions that must be true when there is a discriminator:
+    # 1. Every type in the anyOf/oneOf list must be a $ref to a named schema, such as
+    #    #/components/schemas/X, rather than an inline schema. This is important because
+    #    we may need to use the schema's simple name (X).
+    # 2. There must be a propertyName, representing a property that exists in every
+    #    schema in that list (although we can't currently enforce the latter condition,
+    #    because those properties haven't been parsed yet at this point.)
+    #
+    # There *may* also be a mapping of lookup values (the possible values of the property)
+    # to schemas. Schemas can be referenced either by a full path or a name:
+    #      mapping:
+    #        value_for_a: "#/components/schemas/ModelA"
+    #        value_for_b: ModelB   # equivalent to "#/components/schemas/ModelB"
+    # 
+    # For any type that isn't specified in the mapping (or if the whole mapping is omitted)
+    # the default lookup value for each schema is the same as the schema name. So this--
+    #      mapping:
+    #        value_for_a: "#/components/schemas/ModelA"
+    # --is exactly equivalent to this:
+    #    discriminator:
+    #      propertyName: modelType
+    #      mapping:
+    #        value_for_a: "#/components/schemas/ModelA"
+    #        ModelB: "#/components/schemas/ModelB"
+
+    def _get_model_name(model: ModelProperty) -> str | None:
+        return get_reference_simple_name(model.ref_path) if model.ref_path else None
 
     model_types_by_name: dict[str, PropertyProtocol] = {}
     for model in subtypes:
@@ -245,59 +283,32 @@ def _find_top_level_model(matching_model: ModelProperty) -> ModelProperty | None
             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:
+        name = _get_model_name(model)
+        if not name:
             return PropertyError(
                 detail="Inline schema declarations are not allowed when using a discriminator",
             )
-        name = top_level_model.name
-        if name.startswith("/components/schemas/"):
-            name = get_reference_simple_name(name)
-        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, PropertyProtocol] = model_types_by_name.copy()
+        model_types_by_name[name] = model
+
+    mapping: dict[str, PropertyProtocol] = OrderedDict()  # use ordered dict for test determinacy
+    unspecified_models = list(model_types_by_name.values())
     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 = get_reference_simple_name(ref_path)
-            if not (lookup_model := model_types_by_name.get(name)):
+            if "/" in model_ref:
+                ref_path = parse_reference_path(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 = get_reference_simple_name(ref_path)
+            else:
+                name = model_ref
+            model = model_types_by_name.get(name)
+            if not model:
                 return PropertyError(
-                    detail=f'Discriminator mapping referred to "{model_ref}" which is not one of the schema variants',
+                    detail=f'Discriminator mapping referred to "{name}" 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
-
+            mapping[discriminator_value] = model
+            unspecified_models.remove(model)
+    for model in unspecified_models:
+        if name := _get_model_name(model):
+            mapping[name] = model
     return DiscriminatorDefinition(property_name=data.propertyName, value_to_model_map=mapping)
-
-
-def _validate_discriminators(
-    discriminators: list[DiscriminatorDefinition],
-) -> PropertyError | None:
-    from .model_property import ModelProperty
-
-    prop_names_values_classes = [
-        (discriminator.property_name, key, cast(ModelProperty, 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/templates/property_templates/union_property.py.jinja

@@ -16,13 +16,13 @@ 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) }}
+    {% for value, model in discriminator.value_to_model_map.items() %}
+    def _parse_{{ loop.index }}(data: object) -> {{ model.get_type_string() }}:
+{{ construct_inner_property(model) | indent(8, True) }}
     {% endfor %}
     _discriminator_mapping = {
     {% for value, model in discriminator.value_to_model_map.items() %}
-        "{{ value }}": _parse_{{ model.python_name }},
+        "{{ value }}": _parse_{{ loop.index }},
     {% endfor %}
     }
     if _parse_fn := _discriminator_mapping.get(_discriminator_value):

tests/test_parser/test_properties/test_union.py

@@ -10,7 +10,7 @@
 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.parser.properties.schemas import Class, ReferencePath
 from openapi_python_client.schema import DataType, ParameterLocation
 from tests.test_parser.test_properties.properties_test_helpers import assert_prop_error
 
@@ -66,6 +66,7 @@ def _make_basic_model(
     )
     assert isinstance(model, ModelProperty)
     if name:
+        model.ref_path = ReferencePath(f"/components/schemas/{name}")
         schemas = evolve(
             schemas, classes_by_reference={**schemas.classes_by_reference, f"/components/schemas/{name}": model}
         )
@@ -404,43 +405,3 @@ def test_discriminator_invalid_inline_schema_variant(config, string_property_fac
         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"')