Improve support for unions

libre-man · libre-man · commit 711573eedb0d · 2020-10-12T22:19:30.000+02:00
diff --git a/openapi_python_client/parser/openapi.py b/openapi_python_client/parser/openapi.py
@@ -236,7 +236,7 @@ class Model:
     is_union: bool = False
 
     @staticmethod
-    def from_data(*, data: oai.Schema, name: str) -> Union[Model, ParseError]:
+    def from_data(*, data: oai.Schema, name: str, ref: Reference = None) -> Union[Model, ParseError]:
         """ A single Model from its OAI data
 
         Args:
@@ -249,7 +249,7 @@ def from_data(*, data: oai.Schema, name: str) -> Union[Model, ParseError]:
         optional_properties: List[Property] = []
         relative_imports: Set[str] = set()
 
-        ref = Reference.from_ref(data.title or name)
+        ref = ref or Reference.from_ref(data.title or name)
 
         inherits = None
         props = data.properties
@@ -261,7 +261,6 @@ def from_data(*, data: oai.Schema, name: str) -> Union[Model, ParseError]:
             relative_imports.add(f"from .{inherits.module_name} import {inherits.class_name}")
             props = data.allOf[1].properties
         elif not props:
-            breakpoint()
             props = {}
             # raise AssertionError('Found empty property!')
 
@@ -302,11 +301,21 @@ def from_data(*, data: oai.Schema, name: str) -> Union[MyUnion]:
         ref = Reference.from_ref(data.title or name)
         name = data.title or name
 
-        return MyUnion(
+        model = MyUnion(
             reference=ref,
-            joins=[Model.from_data(data=opt, name=f'{name}_{idx + 1}') for idx, opt in enumerate(data.anyOf)],
+            joins=[
+                Model.from_data(
+                    data=opt,
+                    name=f'{name}_{idx + 1}',
+                    ref=Reference(
+                        class_name=f'{ref.class_name}{idx + 1}',
+                        module_name=ref.module_name,
+                    ),
+                ) for idx, opt in enumerate(data.anyOf)],
             relative_imports=[]
         )
+        ALL_MODELS[ref] = model
+        return model
 
 @dataclass
 class Schemas:
diff --git a/openapi_python_client/parser/properties.py b/openapi_python_client/parser/properties.py
@@ -309,6 +309,10 @@ def get_type_string(self, no_optional: bool = False) -> str:
             return f'Literal[{value_as_string}]'
         return f'Literal[{value_as_string}, None]'
 
+    @property
+    def repr_value(self) -> str:
+        return repr(self.value)
+
     def get_imports(self, *, prefix: str) -> Set[str]:
         """
         Get a set of import strings that should be included when this property is used somewhere
diff --git a/openapi_python_client/templates/property_templates/literal_property.pyi b/openapi_python_client/templates/property_templates/literal_property.pyi
@@ -1,4 +1,6 @@
 {% macro construct(property, source) %}
+if {{ source }} != {{ property.repr_value }}:
+    raise ValueError('{{ "Wrong value for " + property.python_name + ": "}}' + {{ source }})
 {{ property.python_name }} = {{ source }}
 {% endmacro %}
 
diff --git a/openapi_python_client/templates/property_templates/ref_property.pyi b/openapi_python_client/templates/property_templates/ref_property.pyi
@@ -1,10 +1,31 @@
+{% macro _construct(property, source) %}
+{% set model = property.reference.lookup() %}
+{% if model.is_union %}
+from . import {{ model.joins[0].reference.module_name }}
+
+err = None
+for opt in [{% for submodel in model.joins %}{{ submodel.reference.module_name }}.{{ submodel.reference.class_name }}, {% endfor %}]:
+    try:
+        {{ property.python_name }} = opt.from_dict(cast(Dict[str, Any], {{ source }}))
+    except Exception as exc:
+        err = exc
+    else:
+        break
+else:
+    raise err
+del err
+{% else %}
+{{ property.python_name }} = {{ property.reference.class_name }}.from_dict(cast(Dict[str, Any], {{ source }}))
+{% endif %}
+{% endmacro %}
+
 {% macro construct(property, source) %}
 {% if property.required %}
-{{ property.python_name }} = {{ property.reference.class_name }}.from_dict({{ source }})
+{{ _construct(property, source) }}
 {% else %}
 {{ property.python_name }} = None
 if {{ source }} is not None:
-    {{ property.python_name }} = {{ property.reference.class_name }}.from_dict(cast(Dict[str, Any], {{ source }}))
+    {{ _construct(property, source) | indent(4) }}
 {% endif %}
 {% endmacro %}
 
