from datetime import datetime

from typing import (

Any,

Iterable,

Mapping,

Optional,

Sequence,

Tuple,

Type,

TypeVar,

Union,

overload,

)

from uuid import UUID

import sqlalchemy

from sqlalchemy import (

Column,

ColumnElement,

Extract,

FunctionElement,

FunctionFilter,

Label,

Over,

TypeCoerce,

WithinGroup,

)

from sqlalchemy.orm import InstrumentedAttribute, Mapped

from sqlalchemy.sql._typing import (

_ColumnExpressionArgument,

_ColumnExpressionOrLiteralArgument,

_ColumnExpressionOrStrLabelArgument,

)

from sqlalchemy.sql.elements import (

BinaryExpression,

Case,

Cast,

CollectionAggregate,

ColumnClause,

SQLCoreOperations,

TryCast,

UnaryExpression,

)

from sqlalchemy.sql.expression import Select as _Select

from sqlalchemy.sql.roles import TypedColumnsClauseRole

from sqlalchemy.sql.type_api import TypeEngine

from typing_extensions import Literal, Self

_T = TypeVar("_T")

_TypeEngineArgument = Union[Type[TypeEngine[_T]], TypeEngine[_T]]

# Redefine operatos that would only take a column expresion to also take the (virtual)

# types of Pydantic models, e.g. str instead of only Mapped[str].

def all_(expr: Union[_ColumnExpressionArgument[_T], _T]) -> CollectionAggregate[bool]:

return sqlalchemy.all_(expr) # type: ignore[arg-type]

def and_(

initial_clause: Union[Literal[True], _ColumnExpressionArgument[bool], bool],

*clauses: Union[_ColumnExpressionArgument[bool], bool],

) -> ColumnElement[bool]:

return sqlalchemy.and_(initial_clause, *clauses) # type: ignore[arg-type]

def any_(expr: Union[_ColumnExpressionArgument[_T], _T]) -> CollectionAggregate[bool]:

return sqlalchemy.any_(expr) # type: ignore[arg-type]

def asc(

column: Union[_ColumnExpressionOrStrLabelArgument[_T], _T],

) -> UnaryExpression[_T]:

return sqlalchemy.asc(column) # type: ignore[arg-type]

def collate(

expression: Union[_ColumnExpressionArgument[str], str], collation: str

) -> BinaryExpression[str]:

return sqlalchemy.collate(expression, collation) # type: ignore[arg-type]

def between(

expr: Union[_ColumnExpressionOrLiteralArgument[_T], _T],

lower_bound: Any,

upper_bound: Any,

symmetric: bool = False,

) -> BinaryExpression[bool]:

return sqlalchemy.between(expr, lower_bound, upper_bound, symmetric=symmetric) # type: ignore[arg-type]

def not_(clause: Union[_ColumnExpressionArgument[_T], _T]) -> ColumnElement[_T]:

return sqlalchemy.not_(clause) # type: ignore[arg-type]

def case(

*whens: Union[

Tuple[Union[_ColumnExpressionArgument[bool], bool], Any], Mapping[Any, Any]

],

value: Optional[Any] = None,

else_: Optional[Any] = None,

) -> Case[Any]:

return sqlalchemy.case(*whens, value=value, else_=else_) # type: ignore[arg-type]

def cast(

expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any],

type_: "_TypeEngineArgument[_T]",

) -> Cast[_T]:

return sqlalchemy.cast(expression, type_) # type: ignore[arg-type]

def try_cast(

expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any],

type_: "_TypeEngineArgument[_T]",

) -> TryCast[_T]:

return sqlalchemy.try_cast(expression, type_) # type: ignore[arg-type]

def desc(

column: Union[_ColumnExpressionOrStrLabelArgument[_T], _T],

) -> UnaryExpression[_T]:

return sqlalchemy.desc(column) # type: ignore[arg-type]

def distinct(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]:

return sqlalchemy.distinct(expr) # type: ignore[arg-type]

def bitwise_not(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]:

return sqlalchemy.bitwise_not(expr) # type: ignore[arg-type]

def extract(field: str, expr: Union[_ColumnExpressionArgument[Any], Any]) -> Extract:

return sqlalchemy.extract(field, expr) # type: ignore[arg-type]

def funcfilter(

func: FunctionElement[_T], *criterion: Union[_ColumnExpressionArgument[bool], bool]

) -> FunctionFilter[_T]:

return sqlalchemy.funcfilter(func, *criterion) # type: ignore[arg-type]

def label(

name: str,

element: Union[_ColumnExpressionArgument[_T], _T],

type_: Optional["_TypeEngineArgument[_T]"] = None,

) -> Label[_T]:

return sqlalchemy.label(name, element, type_=type_) # type: ignore[arg-type]

def nulls_first(

column: Union[_ColumnExpressionArgument[_T], _T],

) -> UnaryExpression[_T]:

return sqlalchemy.nulls_first(column) # type: ignore[arg-type]

def nulls_last(column: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]:

return sqlalchemy.nulls_last(column) # type: ignore[arg-type]

def or_( # type: ignore[empty-body]

initial_clause: Union[Literal[False], _ColumnExpressionArgument[bool], bool],

*clauses: Union[_ColumnExpressionArgument[bool], bool],

) -> ColumnElement[bool]:

return sqlalchemy.or_(initial_clause, *clauses) # type: ignore[arg-type]

def over(

element: FunctionElement[_T],

partition_by: Optional[

Union[

Iterable[Union[_ColumnExpressionArgument[Any], Any]],

_ColumnExpressionArgument[Any],

Any,

]

] = None,

order_by: Optional[

Union[

Iterable[Union[_ColumnExpressionArgument[Any], Any]],

_ColumnExpressionArgument[Any],

Any,

]

] = None,

range_: Optional[Tuple[Optional[int], Optional[int]]] = None,

rows: Optional[Tuple[Optional[int], Optional[int]]] = None,

) -> Over[_T]:

return sqlalchemy.over(

element, partition_by=partition_by, order_by=order_by, range_=range_, rows=rows

) # type: ignore[arg-type]

def tuple_(

*clauses: Union[_ColumnExpressionArgument[Any], Any],

types: Optional[Sequence["_TypeEngineArgument[Any]"]] = None,

) -> Tuple[Any, ...]:

return sqlalchemy.tuple_(*clauses, types=types) # type: ignore[return-value]

def type_coerce(

expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any],

type_: "_TypeEngineArgument[_T]",

) -> TypeCoerce[_T]:

return sqlalchemy.type_coerce(expression, type_) # type: ignore[arg-type]

def within_group(

element: FunctionElement[_T], *order_by: Union[_ColumnExpressionArgument[Any], Any]

) -> WithinGroup[_T]:

return sqlalchemy.within_group(element, *order_by) # type: ignore[arg-type]

# Separate this class in SelectBase, Select, and SelectOfScalar so that they can share

# where and having without having type overlap incompatibility in session.exec().

class SelectBase(_Select[Tuple[_T]]):

inherit_cache = True

def where(self, *whereclause: Union[_ColumnExpressionArgument[bool], bool]) -> Self:

"""Return a new `Select` construct with the given expression added to

its `WHERE` clause, joined to the existing clause via `AND`, if any.

"""

return super().where(*whereclause) # type: ignore[arg-type]

def having(self, *having: Union[_ColumnExpressionArgument[bool], bool]) -> Self:

"""Return a new `Select` construct with the given expression added to

its `HAVING` clause, joined to the existing clause via `AND`, if any.

"""

return super().having(*having) # type: ignore[arg-type]

class Select(SelectBase[_T]):

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(SelectBase[_T]):

inherit_cache = True

_TCCA = Union[

TypedColumnsClauseRole[_T],

SQLCoreOperations[_T],

Type[_T],

]

# 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,

)

_T{{ i }} = TypeVar("_T{{ i }}")

{% endfor %}

# Generated TypeVars end

@overload

def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:

...

@overload

def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_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 %}

) -> Select[Tuple[{%for ret in signature[1] %}{{ ret }} {% if not loop.last %}, {% endif %}{% endfor %}]]:

...

{% endfor %}

# Generated overloads end

def select(*entities: Any) -> Union[Select, SelectOfScalar]: # type: ignore

if len(entities) == 1:

return SelectOfScalar(*entities)

return Select(*entities)

def col(column_expression: _T) -> Mapped[_T]:

if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)):

raise RuntimeError(f"Not a SQLAlchemy column: {column_expression}")

return column_expression # type: ignore