kittycad.py/generate/union-type.py.jinja2

from typing import Dict, Any, Union, Type, TypeVar
from typing_extensions import Self
import attr
from ..types import UNSET, Unset

from pydantic import GetCoreSchemaHandler
from pydantic_core import CoreSchema, core_schema

GY = TypeVar("GY", bound="{{name}}")

@attr.s(auto_attribs=True)
class {{name}}:
    {% if description %}
    """{{description}}"""
    {% endif %}
    type: Union[
        {% for type in types %}
        {{type.name}},
        {% endfor %}
    ]

    def __init__(self,
                 type: Union[
                     {% for type in types %}
                     {{type.name}},
                     {% endfor %}
                 ]):
        self.type = type

    def model_dump(self) -> Dict[str, Any]:
        {% for type in types %}{% if loop.first %}
        if isinstance(self.type, {{type.name}}):
            {{type.var0}} : {{type.name}} = self.type
            return {{type.var0}}.model_dump()
        {% else %}elif isinstance(self.type, {{type.name}}):
            {{type.var0}} : {{type.name}} = self.type
            return {{type.var0}}.model_dump()
        {% endif %}{% endfor %}
        raise Exception("Unknown type")

    @classmethod
    def from_dict(cls: Type[GY], d: Dict[str, Any]) -> GY:
        {% for type in types %}{% if loop.first %}
        if d.get("{{type.check}}") == {{type.value}}:
            {{type.var1}} : {{type.name}} = {{type.name}}(**d)
            return cls(type={{type.var1}})
        {% else %}elif d.get("{{type.check}}") == {{type.value}}:
            {{type.var1}} : {{type.name}} = {{type.name}}(**d)
            return cls(type={{type.var1}})
        {% endif %}{% endfor %}
        raise Exception("Unknown type")

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> CoreSchema:
        return core_schema.no_info_after_validator_function(cls, handler(Union[
        {% for type in types %}
        {{type.name}},
        {% endfor %}
    ]))