95 lines
2.4 KiB
Django/Jinja
95 lines
2.4 KiB
Django/Jinja
from datetime import datetime
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Generic,
|
|
Mapping,
|
|
Sequence,
|
|
Tuple,
|
|
Type,
|
|
TypeVar,
|
|
Union,
|
|
overload,
|
|
)
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import Column
|
|
from sqlalchemy.orm import InstrumentedAttribute
|
|
from sqlalchemy.sql.elements import ColumnClause
|
|
from sqlalchemy.sql.expression import Select as _Select
|
|
|
|
_TSelect = TypeVar("_TSelect")
|
|
|
|
class Select(_Select, Generic[_TSelect]):
|
|
inherit_cache = True
|
|
|
|
# This is not comparable to sqlalchemy.sql.selectable.ScalarSelect, that has a different
|
|
# purpose. This is the same as a normal SQLAlchemy Select class where there's only one
|
|
# entity, so the result will be converted to a scalar by default. This way writing
|
|
# for loops on the results will feel natural.
|
|
class SelectOfScalar(_Select, Generic[_TSelect]):
|
|
inherit_cache = True
|
|
|
|
if TYPE_CHECKING: # pragma: no cover
|
|
from ..main import SQLModel
|
|
|
|
# Generated TypeVars start
|
|
|
|
{% for i in range(number_of_types) %}
|
|
_TScalar_{{ i }} = TypeVar(
|
|
"_TScalar_{{ i }}",
|
|
Column, # type: ignore
|
|
Sequence, # type: ignore
|
|
Mapping, # type: ignore
|
|
UUID,
|
|
datetime,
|
|
float,
|
|
int,
|
|
bool,
|
|
bytes,
|
|
str,
|
|
None,
|
|
)
|
|
|
|
_TModel_{{ i }} = TypeVar("_TModel_{{ i }}", bound="SQLModel")
|
|
|
|
{% endfor %}
|
|
|
|
# Generated TypeVars end
|
|
|
|
@overload
|
|
def select(entity_0: _TScalar_0, **kw: Any) -> SelectOfScalar[_TScalar_0]: # type: ignore
|
|
...
|
|
|
|
|
|
@overload
|
|
def select(entity_0: Type[_TModel_0], **kw: Any) -> SelectOfScalar[_TModel_0]: # type: ignore
|
|
...
|
|
|
|
|
|
# Generated overloads start
|
|
|
|
{% for signature in signatures %}
|
|
|
|
@overload
|
|
def select( # type: ignore
|
|
{% for arg in signature[0] %}{{ arg.name }}: {{ arg.annotation }}, {% endfor %}**kw: Any,
|
|
) -> Select[Tuple[{%for ret in signature[1] %}{{ ret }} {% if not loop.last %}, {% endif %}{% endfor %}]]:
|
|
...
|
|
|
|
{% endfor %}
|
|
|
|
# Generated overloads end
|
|
|
|
def select(*entities: Any, **kw: Any) -> Union[Select, SelectOfScalar]: # type: ignore
|
|
if len(entities) == 1:
|
|
return SelectOfScalar._create(*entities, **kw) # type: ignore
|
|
return Select._create(*entities, **kw) # type: ignore
|
|
|
|
|
|
# TODO: add several @overload from Python types to SQLAlchemy equivalents
|
|
def col(column_expression: Any) -> ColumnClause: # type: ignore
|
|
if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)):
|
|
raise RuntimeError(f"Not a SQLAlchemy column: {column_expression}")
|
|
return column_expression
|