Feat/detect discriminator bnch 49688 (#143)

* add discriminator

* initialize discriminator mappings

* getter for discriminator python_name

* remove unused method

* fix discriminator.mapping null case

* more inline discriminator properties

* add discriminator_property.pyi

* move discriminator property over to aurelia

* handle unknowntype for discriminator properties

* walk oneOf as well as anyOf references

* add parsing of toplevel UnionProperty

* support toplevel unionProperty

* add nested parameter to construct macro

Author: straz

openapi_python_client/.flake8

@@ -0,0 +1,3 @@
+[flake8]
+per-file-ignores =
+  parser/properties/__init__.py: E402

openapi_python_client/__init__.py

@@ -15,6 +15,7 @@
 
 from .parser import GeneratorData, import_string_from_reference
 from .parser.errors import GeneratorError
+from .parser.properties import UnionProperty
 from .utils import snake_case
 
 if sys.version_info.minor < 8:  # version did not exist before 3.8, need to use a backport
@@ -46,7 +47,6 @@ class Project:
 
     def __init__(self, *, openapi: GeneratorData, custom_template_path: Optional[Path] = None) -> None:
         self.openapi: GeneratorData = openapi
-
         package_loader = PackageLoader(__package__)
         loader: BaseLoader
         if custom_template_path is not None:
@@ -174,10 +174,18 @@ def _build_models(self) -> None:
         imports = []
 
         model_template = self.env.get_template("model.pyi")
+        union_property_template = self.env.get_template("polymorphic_model.pyi")
+
         for model in self.openapi.models.values():
-            module_path = models_dir / f"{model.reference.module_name}.py"
-            module_path.write_text(model_template.render(model=model))
-            imports.append(import_string_from_reference(model.reference))
+            if isinstance(model, UnionProperty):
+                template = union_property_template
+            else:
+                template = model_template
+
+            module_path = models_dir / f"{model.module_name}.py"
+            module_path.write_text(template.render(model=model))
+            if not isinstance(model, UnionProperty):
+                imports.append(import_string_from_reference(model.reference))
 
         # Generate enums
         str_enum_template = self.env.get_template("str_enum.pyi")

openapi_python_client/parser/properties/__init__.py

@@ -1,3 +1,5 @@
+_property = property  # isort: skip
+
 from itertools import chain
 from typing import Any, ClassVar, Dict, Generic, Iterable, Iterator, List, Optional, Set, Tuple, TypeVar, Union
 
@@ -160,10 +162,11 @@ class UnionProperty(Property):
     """ A property representing a Union (anyOf) of other properties """
 
     inner_properties: List[Property]
+    relative_imports: Set[str] = set()
     template: ClassVar[str] = "union_property.pyi"
     has_properties_without_templates: bool = attr.ib(init=False)
-    discriminator_property: Optional[str]
-    discriminator_mappings: Dict[str, Property]
+    discriminator_property: Optional[str] = None
+    discriminator_mappings: Dict[str, Property] = {}
 
     def __attrs_post_init__(self) -> None:
         super().__attrs_post_init__()
@@ -181,6 +184,14 @@ def _get_inner_type_strings(self, json: bool = False) -> List[str]:
     def get_base_type_string(self, json: bool = False) -> str:
         return f"Union[{', '.join(self._get_inner_type_strings(json=json))}]"
 
+    def resolve_references(self, components, schemas):
+        self.relative_imports.update(self.get_imports(prefix=".."))
+        return schemas
+
+    @_property
+    def module_name(self):
+        return self.python_name
+
     def get_type_strings_in_union(
         self, no_optional: bool = False, query_parameter: bool = False, json: bool = False
     ) -> List[str]:
@@ -284,6 +295,14 @@ def build_model_property(
             Used to infer the type name if a `title` property is not available.
         schemas: Existing Schemas which have already been processed (to check name conflicts)
     """
+    if data.anyOf or data.oneOf:
+        prop, schemas = build_union_property(
+            data=data, name=name, required=required, schemas=schemas, parent_name=parent_name
+        )
+        if not isinstance(prop, PropertyError):
+            schemas = attr.evolve(schemas, models={**schemas.models, prop.name: prop})
+        return prop, schemas
+
     required_set = set(data.required or [])
     required_properties: List[Property] = []
     optional_properties: List[Property] = []
@@ -317,6 +336,11 @@ def build_model_property(
             optional_properties.append(prop)
         relative_imports.update(prop.get_imports(prefix=".."))
 
+    discriminator_mappings: Dict[str, Property] = {}
+    if data.discriminator is not None:
+        for k, v in (data.discriminator.mapping or {}).items():
+            discriminator_mappings[k] = Reference.from_ref(v)
+
     additional_properties: Union[bool, Property, PropertyError]
     if data.additionalProperties is None:
         additional_properties = True
@@ -347,6 +371,8 @@ def build_model_property(
         description=data.description or "",
         default=None,
         nullable=data.nullable,
+        discriminator_property=data.discriminator.propertyName if data.discriminator else None,
+        discriminator_mappings=discriminator_mappings,
         required=required,
         name=name,
         additional_properties=additional_properties,
@@ -446,7 +472,7 @@ def build_union_property(
 
     discriminator_mappings: Dict[str, Property] = {}
     if data.discriminator is not None:
-        for k, v in (data.discriminator.mapping if data.discriminator else {}).items():
+        for k, v in (data.discriminator.mapping or {}).items():
             ref_class_name = Reference.from_ref(v).class_name
             discriminator_mappings[k] = reference_name_to_subprop[ref_class_name]
 

openapi_python_client/parser/properties/model_property.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, ClassVar, Dict, List, Set, Union
+from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Set, Union
 
 import attr
 
@@ -22,6 +22,9 @@ class ModelProperty(Property):
     references: List[oai.Reference]
     required_properties: List[Property]
     optional_properties: List[Property]
+    discriminator_property: Optional[str]
+    discriminator_mappings: Dict[str, Property]
+
     description: str
     relative_imports: Set[str]
     additional_properties: Union[bool, Property]
@@ -30,6 +33,10 @@ class ModelProperty(Property):
     template: ClassVar[str] = "model_property.pyi"
     json_is_dict: ClassVar[bool] = True
 
+    @property
+    def module_name(self):
+        return self.reference.module_name
+
     def resolve_references(
         self, components: Dict[str, Union[oai.Reference, oai.Schema]], schemas: Schemas
     ) -> Union[Schemas, PropertyError]:
@@ -44,7 +51,7 @@ def resolve_references(
             assert isinstance(referenced_prop, oai.Schema)
             for p, val in (referenced_prop.properties or {}).items():
                 props[p] = (val, source_name)
-            for sub_prop in referenced_prop.allOf or []:
+            for sub_prop in referenced_prop.allOf or referenced_prop.anyOf or referenced_prop.oneOf or []:
                 if isinstance(sub_prop, oai.Reference):
                     self.references.append(sub_prop)
                 else:
@@ -71,9 +78,17 @@ def resolve_references(
                 self.optional_properties.append(prop)
             self.relative_imports.update(prop.get_imports(prefix=".."))
 
+        for _, value in self.discriminator_mappings.items():
+            self.relative_imports.add(f"from ..models.{value.module_name} import {value.class_name}")
+
         return schemas
 
     def get_base_type_string(self) -> str:
+        if getattr(self, "discriminator_mappings", None):
+            discriminator_types = ", ".join(
+                [ref.class_name for ref in self.discriminator_mappings.values()] + ["UnknownType"]
+            )
+            return f"Union[{discriminator_types}]"
         return self.reference.class_name
 
     def get_imports(self, *, prefix: str) -> Set[str]:

openapi_python_client/templates/property_templates/date_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value="None") %}
+{% macro construct(property, source, initial_value="None", nested=False) %}
 {% if property.required and not property.nullable %}
 {{ property.python_name }} = isoparse({{ source }}).date()
 {% else %}

openapi_python_client/templates/property_templates/datetime_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value="None") %}
+{% macro construct(property, source, initial_value="None", nested=False) %}
 {% if property.required %}
 {% if property.nullable %}
 {{ property.python_name }} = {{ source }}

openapi_python_client/templates/property_templates/dict_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value="None") %}
+{% macro construct(property, source, initial_value="None", nested=False) %}
 {% if property.required %}
 {{ property.python_name }} = {{ source }}
 {% else %}

openapi_python_client/templates/property_templates/enum_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value="None") %}
+{% macro construct(property, source, initial_value="None", nested=False) %}
 {% if property.required %}
 {{ property.python_name }} = {{ property.reference.class_name }}({{ source }})
 {% else %}

openapi_python_client/templates/property_templates/file_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value=None) %}
+{% macro construct(property, source, initial_value=None, nested=False) %}
 {{ property.python_name }} = File(
      payload = BytesIO({{ source }})
 )

openapi_python_client/templates/property_templates/list_property.pyi

@@ -1,4 +1,4 @@
-{% macro construct(property, source, initial_value="[]") %}
+{% macro construct(property, source, initial_value="[]", nested=False) %}
 {% set inner_property = property.inner_property %}
 {% if inner_property.template %}
 {% set inner_source = inner_property.python_name + "_data" %}