Typed queries 2 (reduced) (#149)

This PR includes a reduced version of the changes proposed in #81 .

I removed the plugin method hook and its tests to reduce the scope of the PR, to make it easier to manage and merge, as just the type annotations are a great improvement.

Author: tiangolo

sqlalchemy-stubs/orm/query.pyi

@@ -1,35 +1,44 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, TypeVar, Generic, List, Iterator
 from . import interfaces
 from .base import InspectionAttr
-from ..sql.selectable import ForUpdateArg
+from ..sql.selectable import ForUpdateArg, Alias, CTE
+from ..sql.elements import Label
+from .session import Session
 
-class Query(object):
-    session: Any = ...
-    def __init__(self, entities, session: Optional[Any] = ...) -> None: ...
+
+_T = TypeVar('_T')
+_Q = TypeVar('_Q', bound="Query")
+
+
+class Query(Generic[_T]):
+    session: Session = ...
+    def __init__(self, entities, session: Optional[Session] = ...) -> None: ...
+
+    # TODO: is "statement" always of type sqlalchemy.sql.selectable.Select ?
     @property
     def statement(self): ...
-    def subquery(self, name: Optional[Any] = ..., with_labels: bool = ..., reduce_columns: bool = ...): ...
-    def cte(self, name: Optional[Any] = ..., recursive: bool = ...): ...
-    def label(self, name): ...
+    def subquery(self, name: Optional[str] = ..., with_labels: bool = ..., reduce_columns: bool = ...) -> Alias: ...
+    def cte(self, name: Optional[str] = ..., recursive: bool = ...) -> CTE: ...
+    def label(self, name: str) -> Label: ...
     def as_scalar(self): ...
     @property
     def selectable(self): ...
     def __clause_element__(self): ...
-    def enable_eagerloads(self, value): ...
-    def with_labels(self): ...
-    def enable_assertions(self, value): ...
+    def enable_eagerloads(self: _Q, value: bool) -> _Q: ...
+    def with_labels(self: _Q) -> _Q: ...
+    def enable_assertions(self: _Q, value: bool) -> _Q: ...
     @property
     def whereclause(self): ...
     def with_polymorphic(self, cls_or_mappers, selectable: Optional[Any] = ...,
                          polymorphic_on: Optional[Any] = ...): ...
-    def yield_per(self, count): ...
-    def get(self, ident): ...
+    def yield_per(self: _Q, count: int) -> _Q: ...
+    def get(self, ident) -> Optional[_T]: ...
     def correlate(self, *args): ...
-    def autoflush(self, setting): ...
-    def populate_existing(self): ...
+    def autoflush(self: _Q, setting: bool) -> _Q: ...
+    def populate_existing(self: _Q) -> _Q: ...
     def with_parent(self, instance, property: Optional[Any] = ...): ...
     def add_entity(self, entity, alias: Optional[Any] = ...): ...
-    def with_session(self, session): ...
+    def with_session(self: _Q, session: Optional[Session]) -> _Q: ...
     def from_self(self, *entities): ...
     def values(self, *columns): ...
     def value(self, column): ...
@@ -42,14 +51,14 @@ class Query(object):
     def with_statement_hint(self, text, dialect_name: str = ...): ...
     def execution_options(self, **kwargs): ...
     def with_lockmode(self, mode): ...
-    def with_for_update(self, read: bool = ..., nowait: bool = ..., of: Optional[Any] = ...,
-                        skip_locked: bool = ..., key_share: bool = ...): ...
-    def params(self, *args, **kwargs): ...
-    def filter(self, *criterion): ...
-    def filter_by(self, **kwargs): ...
-    def order_by(self, *criterion): ...
-    def group_by(self, *criterion): ...
-    def having(self, criterion): ...
+    def with_for_update(self: _Q, read: bool = ..., nowait: bool = ..., of: Optional[Any] = ...,
+                        skip_locked: bool = ..., key_share: bool = ...) -> _Q: ...
+    def params(self: _Q, *args, **kwargs) -> _Q: ...
+    def filter(self: _Q, *criterion) -> _Q: ...
+    def filter_by(self: _Q, **kwargs) -> _Q: ...
+    def order_by(self: _Q, *criterion) -> _Q: ...
+    def group_by(self: _Q, *criterion) -> _Q: ...
+    def having(self: _Q, criterion) -> _Q: ...
     def union(self, *q): ...
     def union_all(self, *q): ...
     def intersect(self, *q): ...
@@ -62,26 +71,26 @@ class Query(object):
     def select_from(self, *from_obj): ...
     def select_entity_from(self, from_obj): ...
     def __getitem__(self, item): ...
-    def slice(self, start, stop): ...
-    def limit(self, limit): ...
-    def offset(self, offset): ...
+    def slice(self: _Q, start: Optional[int], stop: Optional[int]) -> _Q: ...
+    def limit(self: _Q, limit: Optional[int]) -> _Q: ...
+    def offset(self: _Q, offset: Optional[int]) -> _Q: ...
     def distinct(self, *criterion): ...
     def prefix_with(self, *prefixes): ...
     def suffix_with(self, *suffixes): ...
-    def all(self): ...
+    def all(self) -> List[_T]: ...
     def from_statement(self, statement): ...
-    def first(self): ...
-    def one_or_none(self): ...
-    def one(self): ...
+    def first(self) -> Optional[_T]: ...
+    def one_or_none(self) -> Optional[_T]: ...
+    def one(self) -> _T: ...
     def scalar(self): ...
-    def __iter__(self): ...
+    def __iter__(self) -> Iterator[_T]: ...
     @property
     def column_descriptions(self): ...
     def instances(self, cursor, __context: Optional[Any] = ...): ...
     def merge_result(self, iterator, load: bool = ...): ...
     def exists(self): ...
-    def count(self): ...
-    def delete(self, synchronize_session: Union[bool, str] = ...): ...
+    def count(self) -> int: ...
+    def delete(self, synchronize_session: Union[bool, str] = ...) -> int: ...
     def update(self, values, synchronize_session: Union[bool, str] = ..., update_args: Optional[Any] = ...): ...
 
 class LockmodeArg(ForUpdateArg):

sqlalchemy-stubs/orm/session.pyi

@@ -1,5 +1,7 @@
 from typing import Any, Optional
 
+from sqlalchemy.orm.query import Query
+
 class _SessionClassMethods(object):
     @classmethod
     def close_all(cls): ...
@@ -60,7 +62,7 @@ class Session(_SessionClassMethods):
     def bind_mapper(self, mapper, bind): ...
     def bind_table(self, table, bind): ...
     def get_bind(self, mapper: Optional[Any] = ..., clause: Optional[Any] = ...): ...
-    def query(self, *entities, **kwargs): ...
+    def query(self, *entities, **kwargs) -> Query[Any]: ...
     @property
     def no_autoflush(self): ...
     def refresh(self, instance, attribute_names: Optional[Any] = ..., lockmode: Optional[Any] = ...): ...