120 lines
3.1 KiB
Django/Jinja
120 lines
3.1 KiB
Django/Jinja
import sys
|
|
from datetime import datetime
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Generic,
|
|
Mapping,
|
|
Sequence,
|
|
Tuple,
|
|
Type,
|
|
TypeVar,
|
|
Union,
|
|
cast,
|
|
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")
|
|
|
|
# Workaround Generics incompatibility in Python 3.6
|
|
# Ref: https://github.com/python/typing/issues/449#issuecomment-316061322
|
|
if sys.version_info.minor >= 7:
|
|
|
|
class Select(_Select, Generic[_TSelect]):
|
|
pass
|
|
|
|
# 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]):
|
|
pass
|
|
|
|
|
|
else:
|
|
from typing import GenericMeta # type: ignore
|
|
|
|
class GenericSelectMeta(GenericMeta, _Select.__class__): # type: ignore
|
|
pass
|
|
|
|
class _Py36Select(_Select, Generic[_TSelect], metaclass=GenericSelectMeta):
|
|
pass
|
|
|
|
class _Py36SelectOfScalar(_Select, Generic[_TSelect], metaclass=GenericSelectMeta):
|
|
pass
|
|
|
|
# Cast them for editors to work correctly, from several tricks tried, this works
|
|
# for both VS Code and PyCharm
|
|
Select = cast("Select", _Py36Select) # type: ignore
|
|
SelectOfScalar = cast("SelectOfScalar", _Py36SelectOfScalar) # type: ignore
|
|
|
|
|
|
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,
|
|
Sequence,
|
|
Mapping,
|
|
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]:
|
|
if len(entities) == 1:
|
|
return SelectOfScalar._create(*entities, **kw) # type: ignore
|
|
return Select._create(*entities, **kw)
|
|
|
|
|
|
# TODO: add several @overload from Python types to SQLAlchemy equivalents
|
|
def col(column_expression: Any) -> ColumnClause:
|
|
if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)):
|
|
raise RuntimeError(f"Not a SQLAlchemy column: {column_expression}")
|
|
return column_expression
|