codekingpro/portable-devtools

1from __future__ import annotations2 3import re4import typing as t5 6import sqlalchemy as sa7import sqlalchemy.orm as sa_orm8 9from .query import Query10 11if t.TYPE_CHECKING:12    from .extension import SQLAlchemy13 14 15class _QueryProperty:16    """A class property that creates a query object for a model.17 18    :meta private:19    """20 21    def __get__(self, obj: Model | None, cls: type[Model]) -> Query:22        return cls.query_class(23            cls, session=cls.__fsa__.session()  # type: ignore[arg-type]24        )25 26 27class Model:28    """The base class of the :attr:`.SQLAlchemy.Model` declarative model class.29 30    To define models, subclass :attr:`db.Model <.SQLAlchemy.Model>`, not this. To31    customize ``db.Model``, subclass this and pass it as ``model_class`` to32    :class:`.SQLAlchemy`. To customize ``db.Model`` at the metaclass level, pass an33    already created declarative model class as ``model_class``.34    """35 36    __fsa__: t.ClassVar[SQLAlchemy]37    """Internal reference to the extension object.38 39    :meta private:40    """41 42    query_class: t.ClassVar[type[Query]] = Query43    """Query class used by :attr:`query`. Defaults to :attr:`.SQLAlchemy.Query`, which44    defaults to :class:`.Query`.45    """46 47    query: t.ClassVar[Query] = _QueryProperty()  # type: ignore[assignment]48    """A SQLAlchemy query for a model. Equivalent to ``db.session.query(Model)``. Can be49    customized per-model by overriding :attr:`query_class`.50 51    .. warning::52        The query interface is considered legacy in SQLAlchemy. Prefer using53        ``session.execute(select())`` instead.54    """55 56    def __repr__(self) -> str:57        state = sa.inspect(self)58        assert state is not None59 60        if state.transient:61            pk = f"(transient {id(self)})"62        elif state.pending:63            pk = f"(pending {id(self)})"64        else:65            pk = ", ".join(map(str, state.identity))66 67        return f"<{type(self).__name__} {pk}>"68 69 70class BindMetaMixin(type):71    """Metaclass mixin that sets a model's ``metadata`` based on its ``__bind_key__``.72 73    If the model sets ``metadata`` or ``__table__`` directly, ``__bind_key__`` is74    ignored. If the ``metadata`` is the same as the parent model, it will not be set75    directly on the child model.76    """77 78    __fsa__: SQLAlchemy79    metadata: sa.MetaData80 81    def __init__(82        cls, name: str, bases: tuple[type, ...], d: dict[str, t.Any], **kwargs: t.Any83    ) -> None:84        if not ("metadata" in cls.__dict__ or "__table__" in cls.__dict__):85            bind_key = getattr(cls, "__bind_key__", None)86            parent_metadata = getattr(cls, "metadata", None)87            metadata = cls.__fsa__._make_metadata(bind_key)88 89            if metadata is not parent_metadata:90                cls.metadata = metadata91 92        super().__init__(name, bases, d, **kwargs)93 94 95class BindMixin:96    """DeclarativeBase mixin to set a model's ``metadata`` based on ``__bind_key__``.97 98    If no ``__bind_key__`` is specified, the model will use the default metadata99    provided by ``DeclarativeBase`` or ``DeclarativeBaseNoMeta``.100    If the model doesn't set ``metadata`` or ``__table__`` directly101    and does set ``__bind_key__``, the model will use the metadata102    for the specified bind key.103    If the ``metadata`` is the same as the parent model, it will not be set104    directly on the child model.105 106    .. versionchanged:: 3.1.0107    """108 109    __fsa__: SQLAlchemy110    metadata: sa.MetaData111 112    @classmethod113    def __init_subclass__(cls: t.Type[BindMixin], **kwargs: t.Dict[str, t.Any]) -> None:114        if not ("metadata" in cls.__dict__ or "__table__" in cls.__dict__) and hasattr(115            cls, "__bind_key__"116        ):117            bind_key = getattr(cls, "__bind_key__", None)118            parent_metadata = getattr(cls, "metadata", None)119            metadata = cls.__fsa__._make_metadata(bind_key)120 121            if metadata is not parent_metadata:122                cls.metadata = metadata123 124        super().__init_subclass__(**kwargs)125 126 127class NameMetaMixin(type):128    """Metaclass mixin that sets a model's ``__tablename__`` by converting the129    ``CamelCase`` class name to ``snake_case``. A name is set for non-abstract models130    that do not otherwise define ``__tablename__``. If a model does not define a primary131    key, it will not generate a name or ``__table__``, for single-table inheritance.132    """133 134    metadata: sa.MetaData135    __tablename__: str136    __table__: sa.Table137 138    def __init__(139        cls, name: str, bases: tuple[type, ...], d: dict[str, t.Any], **kwargs: t.Any140    ) -> None:141        if should_set_tablename(cls):142            cls.__tablename__ = camel_to_snake_case(cls.__name__)143 144        super().__init__(name, bases, d, **kwargs)145 146        # __table_cls__ has run. If no table was created, use the parent table.147        if (148            "__tablename__" not in cls.__dict__149            and "__table__" in cls.__dict__150            and cls.__dict__["__table__"] is None151        ):152            del cls.__table__153 154    def __table_cls__(cls, *args: t.Any, **kwargs: t.Any) -> sa.Table | None:155        """This is called by SQLAlchemy during mapper setup. It determines the final156        table object that the model will use.157 158        If no primary key is found, that indicates single-table inheritance, so no table159        will be created and ``__tablename__`` will be unset.160        """161        schema = kwargs.get("schema")162 163        if schema is None:164            key = args[0]165        else:166            key = f"{schema}.{args[0]}"167 168        # Check if a table with this name already exists. Allows reflected tables to be169        # applied to models by name.170        if key in cls.metadata.tables:171            return sa.Table(*args, **kwargs)172 173        # If a primary key is found, create a table for joined-table inheritance.174        for arg in args:175            if (isinstance(arg, sa.Column) and arg.primary_key) or isinstance(176                arg, sa.PrimaryKeyConstraint177            ):178                return sa.Table(*args, **kwargs)179 180        # If no base classes define a table, return one that's missing a primary key181        # so SQLAlchemy shows the correct error.182        for base in cls.__mro__[1:-1]:183            if "__table__" in base.__dict__:184                break185        else:186            return sa.Table(*args, **kwargs)187 188        # Single-table inheritance, use the parent table name. __init__ will unset189        # __table__ based on this.190        if "__tablename__" in cls.__dict__:191            del cls.__tablename__192 193        return None194 195 196class NameMixin:197    """DeclarativeBase mixin that sets a model's ``__tablename__`` by converting the198    ``CamelCase`` class name to ``snake_case``. A name is set for non-abstract models199    that do not otherwise define ``__tablename__``. If a model does not define a primary200    key, it will not generate a name or ``__table__``, for single-table inheritance.201 202    .. versionchanged:: 3.1.0203    """204 205    metadata: sa.MetaData206    __tablename__: str207    __table__: sa.Table208 209    @classmethod210    def __init_subclass__(cls: t.Type[NameMixin], **kwargs: t.Dict[str, t.Any]) -> None:211        if should_set_tablename(cls):212            cls.__tablename__ = camel_to_snake_case(cls.__name__)213 214        super().__init_subclass__(**kwargs)215 216        # __table_cls__ has run. If no table was created, use the parent table.217        if (218            "__tablename__" not in cls.__dict__219            and "__table__" in cls.__dict__220            and cls.__dict__["__table__"] is None221        ):222            del cls.__table__223 224    @classmethod225    def __table_cls__(cls, *args: t.Any, **kwargs: t.Any) -> sa.Table | None:226        """This is called by SQLAlchemy during mapper setup. It determines the final227        table object that the model will use.228 229        If no primary key is found, that indicates single-table inheritance, so no table230        will be created and ``__tablename__`` will be unset.231        """232        schema = kwargs.get("schema")233 234        if schema is None:235            key = args[0]236        else:237            key = f"{schema}.{args[0]}"238 239        # Check if a table with this name already exists. Allows reflected tables to be240        # applied to models by name.241        if key in cls.metadata.tables:242            return sa.Table(*args, **kwargs)243 244        # If a primary key is found, create a table for joined-table inheritance.245        for arg in args:246            if (isinstance(arg, sa.Column) and arg.primary_key) or isinstance(247                arg, sa.PrimaryKeyConstraint248            ):249                return sa.Table(*args, **kwargs)250 251        # If no base classes define a table, return one that's missing a primary key252        # so SQLAlchemy shows the correct error.253        for base in cls.__mro__[1:-1]:254            if "__table__" in base.__dict__:255                break256        else:257            return sa.Table(*args, **kwargs)258 259        # Single-table inheritance, use the parent table name. __init__ will unset260        # __table__ based on this.261        if "__tablename__" in cls.__dict__:262            del cls.__tablename__263 264        return None265 266 267def should_set_tablename(cls: type) -> bool:268    """Determine whether ``__tablename__`` should be generated for a model.269 270    -   If no class in the MRO sets a name, one should be generated.271    -   If a declared attr is found, it should be used instead.272    -   If a name is found, it should be used if the class is a mixin, otherwise one273        should be generated.274    -   Abstract models should not have one generated.275 276    Later, ``__table_cls__`` will determine if the model looks like single or277    joined-table inheritance. If no primary key is found, the name will be unset.278    """279    if (280        cls.__dict__.get("__abstract__", False)281        or (282            not issubclass(cls, (sa_orm.DeclarativeBase, sa_orm.DeclarativeBaseNoMeta))283            and not any(isinstance(b, sa_orm.DeclarativeMeta) for b in cls.__mro__[1:])284        )285        or any(286            (b is sa_orm.DeclarativeBase or b is sa_orm.DeclarativeBaseNoMeta)287            for b in cls.__bases__288        )289    ):290        return False291 292    for base in cls.__mro__:293        if "__tablename__" not in base.__dict__:294            continue295 296        if isinstance(base.__dict__["__tablename__"], sa_orm.declared_attr):297            return False298 299        return not (300            base is cls301            or base.__dict__.get("__abstract__", False)302            or not (303                # SQLAlchemy 1.x304                isinstance(base, sa_orm.DeclarativeMeta)305                # 2.x: DeclarativeBas uses this as metaclass306                or isinstance(base, sa_orm.decl_api.DeclarativeAttributeIntercept)307                # 2.x: DeclarativeBaseNoMeta doesn't use a metaclass308                or issubclass(base, sa_orm.DeclarativeBaseNoMeta)309            )310        )311 312    return True313 314 315def camel_to_snake_case(name: str) -> str:316    """Convert a ``CamelCase`` name to ``snake_case``."""317    name = re.sub(r"((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))", r"_\1", name)318    return name.lower().lstrip("_")319 320 321class DefaultMeta(BindMetaMixin, NameMetaMixin, sa_orm.DeclarativeMeta):322    """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and323    ``__tablename__`` support.324    """325 326 327class DefaultMetaNoName(BindMetaMixin, sa_orm.DeclarativeMeta):328    """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and329    ``__tablename__`` support.330    """331