Improve orm.query (#70)

* Improve `orm.query`

* Updates based on PR feedback

Author: bryanforbes

sqlalchemy-stubs/orm/query.pyi

@@ -1,18 +1,41 @@
 from typing import Any
+from typing import Generic
+from typing import Iterable
+from typing import Iterator
 from typing import List
 from typing import Optional
+from typing import Tuple
+from typing import TypeVar
+from typing import Union
+
+from typing_extensions import Protocol
 
 from . import interfaces
-from .context import QueryContext as QueryContext
-from .util import aliased as aliased
+from .session import _SessionProtocol
+from .state import InstanceState
 from ..sql.annotation import SupportsCloneAnnotations
 from ..sql.base import Executable
+from ..sql.elements import BooleanClauseList
+from ..sql.elements import GroupedElement
+from ..sql.elements import Label
 from ..sql.selectable import _SelectFromElements
-from ..sql.selectable import GroupedElement
+from ..sql.selectable import Alias
+from ..sql.selectable import CTE
+from ..sql.selectable import Exists
 from ..sql.selectable import HasHints
 from ..sql.selectable import HasPrefixes
 from ..sql.selectable import HasSuffixes
+from ..sql.selectable import ScalarSelect
+from ..sql.selectable import Select
 from ..sql.selectable import SelectBase
+from ..sql.selectable import SelectStatementGrouping
+from ..util.langhelpers import _symbol
+
+_T = TypeVar("_T")
+_TQuery = TypeVar("_TQuery", bound=Query[Any])
+
+class _WithTransformationFn(Protocol[_TQuery]):
+    def __call__(self, __query: _TQuery) -> _TQuery: ...
 
 class Query(
     _SelectFromElements,
@@ -21,144 +44,163 @@ class Query(
     HasSuffixes,
     HasHints,
     Executable,
+    Generic[_T],
 ):
     load_options: Any = ...
-    session: Any = ...
+    session: Optional[_SessionProtocol] = ...
     def __init__(
-        self, entities: Any, session: Optional[Any] = ...
+        self,
+        entities: Iterable[Any],
+        session: Optional[_SessionProtocol] = ...,
     ) -> None: ...
     @property
-    def statement(self): ...
+    def statement(self) -> Union[Select, FromStatement]: ...
     def subquery(
         self,
-        name: Optional[Any] = ...,
+        name: Optional[str] = ...,
         with_labels: bool = ...,
         reduce_columns: bool = ...,
-    ): ...
-    def cte(self, name: Optional[Any] = ..., recursive: bool = ...): ...
-    def label(self, name: Any): ...
-    def as_scalar(self): ...
-    def scalar_subquery(self): ...
-    def __clause_element__(self): ...
-    def only_return_tuples(self, value: Any) -> None: ...
+    ) -> Alias: ...
+    def cte(self, name: Optional[str] = ..., recursive: bool = ...) -> CTE: ...
+    def label(self, name: str) -> Label[Any]: ...
+    def as_scalar(self) -> ScalarSelect: ...
+    def scalar_subquery(self) -> ScalarSelect: ...
+    @property
+    def selectable(self) -> Union[Select, FromStatement]: ...
+    def __clause_element__(self) -> Union[Select, FromStatement]: ...
+    def only_return_tuples(self: _TQuery, value: bool) -> _TQuery: ...
     @property
-    def is_single_entity(self): ...
-    def enable_eagerloads(self, value: Any) -> None: ...
-    def with_labels(self): ...
-    apply_labels: Any = ...
+    def is_single_entity(self) -> bool: ...
+    def enable_eagerloads(self: _TQuery, value: bool) -> _TQuery: ...
+    def with_labels(self: _TQuery) -> _TQuery: ...
+    def apply_labels(self: _TQuery) -> _TQuery: ...
     @property
-    def get_label_style(self): ...
-    def set_label_style(self, style: Any): ...
-    def enable_assertions(self, value: Any) -> None: ...
+    def get_label_style(self) -> _symbol: ...
+    def set_label_style(self: _TQuery, style: _symbol) -> _TQuery: ...
+    def enable_assertions(self: _TQuery, value: bool) -> _TQuery: ...
     @property
-    def whereclause(self): ...
+    def whereclause(self) -> BooleanClauseList[Any]: ...
     def with_polymorphic(
-        self,
+        self: _TQuery,
         cls_or_mappers: Any,
         selectable: Optional[Any] = ...,
         polymorphic_on: Optional[Any] = ...,
-    ) -> None: ...
-    def yield_per(self, count: Any) -> None: ...
-    def get(self, ident: Any): ...
+    ) -> _TQuery: ...
+    def yield_per(self: _TQuery, count: int) -> _TQuery: ...
+    def get(self, ident: Any) -> Optional[_T]: ...
     @property
-    def lazy_loaded_from(self): ...
-    def correlate(self, *fromclauses: Any) -> None: ...
-    def autoflush(self, setting: Any) -> None: ...
-    def populate_existing(self) -> None: ...
+    def lazy_loaded_from(self) -> InstanceState: ...
+    def correlate(self: _TQuery, *fromclauses: Any) -> _TQuery: ...
+    def autoflush(self: _TQuery, setting: Any) -> _TQuery: ...
+    def populate_existing(self: _TQuery) -> _TQuery: ...
     def with_parent(
-        self,
+        self: _TQuery,
         instance: Any,
         property: Optional[Any] = ...,
         from_entity: Optional[Any] = ...,
-    ): ...
-    def add_entity(self, entity: Any, alias: Optional[Any] = ...) -> None: ...
-    def with_session(self, session: Any) -> None: ...
-    def from_self(self, *entities: Any): ...
-    def values(self, *columns: Any): ...
-    def value(self, column: Any): ...
-    def with_entities(self, *entities: Any) -> None: ...
-    def add_columns(self, *column: Any) -> None: ...
-    def add_column(self, column: Any): ...
-    def options(self, *args: Any) -> None: ...
-    def with_transformation(self, fn: Any): ...
-    def get_execution_options(self): ...
-    def execution_options(self, **kwargs: Any) -> None: ...
+    ) -> _TQuery: ...
+    def add_entity(
+        self: _TQuery, entity: Any, alias: Optional[Any] = ...
+    ) -> _TQuery: ...
+    def with_session(
+        self: _TQuery, session: Optional[_SessionProtocol]
+    ) -> _TQuery: ...
+    def from_self(self: _TQuery, *entities: Any) -> _TQuery: ...
+    def values(self, *columns: Any) -> Iterator[Tuple[Any, ...]]: ...
+    def value(self, column: Any) -> Optional[Any]: ...
+    def with_entities(self: _TQuery, *entities: Any) -> _TQuery: ...
+    def add_columns(self: _TQuery, *column: Any) -> _TQuery: ...
+    def add_column(self: _TQuery, column: Any) -> _TQuery: ...
+    def options(self: _TQuery, *args: Any) -> _TQuery: ...
+    def with_transformation(
+        self: _TQuery, fn: _WithTransformationFn[_TQuery]
+    ) -> _TQuery: ...
+    def get_execution_options(self) -> Any: ...
+    def execution_options(self: _TQuery, **kwargs: Any) -> _TQuery: ...
     def with_for_update(
-        self,
+        self: _TQuery,
         read: bool = ...,
         nowait: bool = ...,
         of: Optional[Any] = ...,
         skip_locked: bool = ...,
         key_share: bool = ...,
-    ) -> None: ...
-    def params(self, *args: Any, **kwargs: Any) -> None: ...
-    def where(self, *criterion: Any): ...
-    def filter(self, *criterion: Any) -> None: ...
-    def filter_by(self, **kwargs: Any): ...
-    def order_by(self, *clauses: Any) -> None: ...
-    def group_by(self, *clauses: Any) -> None: ...
-    def having(self, criterion: Any) -> None: ...
-    def union(self, *q: Any): ...
-    def union_all(self, *q: Any): ...
-    def intersect(self, *q: Any): ...
-    def intersect_all(self, *q: Any): ...
-    def except_(self, *q: Any): ...
-    def except_all(self, *q: Any): ...
-    def join(self, target: Any, *props: Any, **kwargs: Any) -> None: ...
-    def outerjoin(self, target: Any, *props: Any, **kwargs: Any): ...
-    def reset_joinpoint(self) -> None: ...
-    def select_from(self, *from_obj: Any) -> None: ...
-    def select_entity_from(self, from_obj: Any) -> None: ...
-    def __getitem__(self, item: Any): ...
-    def slice(self, start: Any, stop: Any) -> None: ...
-    def limit(self, limit: Any) -> None: ...
-    def offset(self, offset: Any) -> None: ...
-    def distinct(self, *expr: Any) -> None: ...
-    def all(self) -> List[Any]: ...
-    def from_statement(self, statement: Any) -> None: ...
-    def first(self): ...
-    def one_or_none(self): ...
-    def one(self): ...
-    def scalar(self): ...
-    def __iter__(self) -> Any: ...
+    ) -> _TQuery: ...
+    def params(self: _TQuery, *args: Any, **kwargs: Any) -> _TQuery: ...
+    def where(self: _TQuery, *criterion: Any) -> _TQuery: ...
+    def filter(self: _TQuery, *criterion: Any) -> _TQuery: ...
+    def filter_by(self: _TQuery, **kwargs: Any) -> _TQuery: ...
+    def order_by(self: _TQuery, *clauses: Any) -> _TQuery: ...
+    def group_by(self: _TQuery, *clauses: Any) -> _TQuery: ...
+    def having(self: _TQuery, criterion: Any) -> _TQuery: ...
+    def union(self: _TQuery, *q: Any) -> _TQuery: ...
+    def union_all(self: _TQuery, *q: Any) -> _TQuery: ...
+    def intersect(self: _TQuery, *q: Any) -> _TQuery: ...
+    def intersect_all(self: _TQuery, *q: Any) -> _TQuery: ...
+    def except_(self: _TQuery, *q: Any) -> _TQuery: ...
+    def except_all(self: _TQuery, *q: Any) -> _TQuery: ...
+    def join(
+        self: _TQuery, target: Any, *props: Any, **kwargs: Any
+    ) -> _TQuery: ...
+    def outerjoin(
+        self: _TQuery, target: Any, *props: Any, **kwargs: Any
+    ) -> _TQuery: ...
+    def reset_joinpoint(self: _TQuery) -> _TQuery: ...
+    def select_from(self: _TQuery, *from_obj: Any) -> _TQuery: ...
+    def select_entity_from(self: _TQuery, from_obj: Any) -> _TQuery: ...
+    def __getitem__(self, item: Any) -> Any: ...
+    def slice(self: _TQuery, start: Any, stop: Any) -> _TQuery: ...
+    def limit(self: _TQuery, limit: Any) -> _TQuery: ...
+    def offset(self: _TQuery, offset: Any) -> _TQuery: ...
+    def distinct(self: _TQuery, *expr: Any) -> _TQuery: ...
+    def all(self) -> List[_T]: ...
+    def from_statement(self: _TQuery, statement: Any) -> _TQuery: ...
+    def first(self) -> Optional[_T]: ...
+    def one_or_none(self) -> Optional[_T]: ...
+    def one(self) -> _T: ...
+    def scalar(self) -> Any: ...  # type: ignore[override]
+    def __iter__(self) -> Iterator[_T]: ...
     @property
-    def column_descriptions(self): ...
-    def instances(self, result_proxy: Any, context: Optional[Any] = ...): ...
-    def merge_result(self, iterator: Any, load: bool = ...): ...
-    def exists(self): ...
-    def count(self): ...
-    def delete(self, synchronize_session: str = ...): ...
+    def column_descriptions(self) -> List[Any]: ...
+    def instances(
+        self, result_proxy: Any, context: Optional[Any] = ...
+    ) -> Any: ...
+    def merge_result(self, iterator: Any, load: bool = ...) -> Any: ...
+    def exists(self) -> Exists: ...
+    def count(self) -> int: ...
+    def delete(self, synchronize_session: str = ...) -> int: ...
     def update(
         self,
         values: Any,
         synchronize_session: str = ...,
         update_args: Optional[Any] = ...,
-    ): ...
+    ) -> int: ...
 
 class FromStatement(GroupedElement, SelectBase, Executable):
     __visit_name__: str = ...
     element: Any = ...
-    def __init__(self, entities: Any, element: Any) -> None: ...
-    def get_label_style(self): ...
-    def set_label_style(self, label_style: Any): ...
-    def get_children(self, **kw: Any) -> None: ...
+    def __init__(self, entities: Iterable[Any], element: Any) -> None: ...
+    def get_label_style(self) -> _symbol: ...
+    def set_label_style(
+        self, label_style: _symbol
+    ) -> SelectStatementGrouping: ...
+    def get_children(self, **kw: Any) -> Iterable[Any]: ...  # type: ignore[override]
 
 class AliasOption(interfaces.LoaderOption):
     def __init__(self, alias: Any) -> None: ...
     def process_compile_state(self, compile_state: Any) -> None: ...
 
 class BulkUD:
-    query: Any = ...
+    query: Query[Any] = ...
     mapper: Any = ...
-    def __init__(self, query: Any) -> None: ...
+    def __init__(self, query: Query[Any]) -> None: ...
     @property
-    def session(self): ...
+    def session(self) -> Optional[_SessionProtocol]: ...
 
 class BulkUpdate(BulkUD):
     values: Any = ...
     update_kwargs: Any = ...
     def __init__(
-        self, query: Any, values: Any, update_kwargs: Any
+        self, query: Query[Any], values: Any, update_kwargs: Any
     ) -> None: ...
 
 class BulkDelete(BulkUD): ...

sqlalchemy-stubs/orm/session.pyi

@@ -39,7 +39,7 @@ _TSessionTransaction = TypeVar(
 )
 
 _TSharedSessionProtocol = TypeVar(
-    "_TSharedSessionProtocol", bound=_SharedSessionProtocol
+    "_TSharedSessionProtocol", bound=_SharedSessionProtocol[Any]
 )
 _TSessionTransactionProtocol = TypeVar(
     "_TSessionTransactionProtocol", bound=_SessionTransactionProtocol
@@ -406,7 +406,9 @@ class Session(_SessionClassMethods):
         binds: Optional[Mapping[Any, Union[Connection, Engine]]] = ...,
         enable_baked_queries: bool = ...,
         info: Optional[Mapping[Any, Any]] = ...,
-        query_cls: Optional[Union[Query, Callable[..., Query]]] = ...,
+        query_cls: Optional[
+            Union[Type[Query[Any]], Callable[..., Query[Any]]]
+        ] = ...,
     ) -> None: ...
     connection_callable: Any = ...
     def __enter__(self: _TSession) -> _TSession: ...
@@ -473,7 +475,7 @@ class Session(_SessionClassMethods):
         _sa_skip_events: Optional[Any] = ...,
         _sa_skip_for_implicit_returning: bool = ...,
     ) -> Union[Connection, Engine]: ...
-    def query(self, *entities: Any, **kwargs: Any) -> Query: ...
+    def query(self, *entities: Any, **kwargs: Any) -> Query[Any]: ...
     @property
     def no_autoflush(self: _TSession) -> ContextManager[_TSession]: ...
     def refresh(