Add one overload to Session.query to type queries of one single item (#…

…151)

Add one overload to `Session.query()` to handle return types for a query of one single item.

This helps with the specific case of querying a single model or column.

For more than one item in the `query()` it seems that another plugin hook is the only option, as the return is not a plain `Tuple` but a `KeyedTuple`.

This is what I found possible continuing the discussion in #149

sqlalchemy-stubs/orm/session.pyi

@@ -1,5 +1,6 @@
-from typing import Any, Optional
+from typing import Any, Optional, Type, TypeVar, Union, overload
 
+from sqlalchemy import Column
 from sqlalchemy.orm.query import Query
 
 class _SessionClassMethods(object):
@@ -26,6 +27,8 @@ class SessionTransaction(object):
     def __enter__(self): ...
     def __exit__(self, type, value, traceback): ...
 
+_T = TypeVar("_T")
+
 class Session(_SessionClassMethods):
     public_methods: Any = ...
     identity_map: Any = ...
@@ -62,7 +65,10 @@ 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) -> Query[Any]: ...
+    @overload
+    def query(self, entity: Union[Type[_T], Column[_T]], **kwargs) -> Query[_T]: ...
+    @overload
+    def query(self, *entities, **kwargs) -> Query: ...
     @property
     def no_autoflush(self): ...
     def refresh(self, instance, attribute_names: Optional[Any] = ..., lockmode: Optional[Any] = ...): ...