type_adapter_pydantic

type_adapter_pydantic.py

class PydanticFromJSON(TypeDecorator):
    """SQLAlchemy type for storing Pydantic models as JSONB.

    Supports single models, lists, and union types via type parameter:
        # Single model
        evidence: Mapped[MyModel] = mapped_column(PydanticFromJSON[MyModel])

        # List of models
        overrides: Mapped[list[MyModel]] = mapped_column(PydanticFromJSON[list[MyModel]])

        # Union types (including enums + models)
        group_by: Mapped[list[EnumType | ModelType]] = mapped_column(
            PydanticFromJSON[list[EnumType | ModelType]]
        )
    """

    impl = JSONB
    cache_ok = False

    # Set by __class_getitem__ when creating parameterized types
    _type_adapter: TypeAdapter

    def __class_getitem__(cls, item):
        """Handle PydanticFromJSON[T] syntax for any supported type T."""
        origin = typing.get_origin(item)

        # Generate descriptive class name from the type
        type_name = getattr(item, "__name__", str(item))

        if origin is list:
            inner_type = typing.get_args(item)[0]
            inner_name = getattr(inner_type, "__name__", str(inner_type))
            new_cls = type(
                f"PydanticFromJSON[list[{inner_name}]]",
                (cls,),
                {"_type_adapter": TypeAdapter(item)},
            )
            return new_cls()

        new_cls = type(
            f"PydanticFromJSON[{type_name}]",
            (cls,),
            {"_type_adapter": TypeAdapter(item)},
        )
        return new_cls()

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        # TypeAdapter handles enums, models, unions, and lists correctly
        return self._type_adapter.dump_python(value, mode="json")

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        # TypeAdapter validates and coerces to the correct types
        return self._type_adapter.validate_python(value)