Improve typing for the 'get' method in session (#132)

* Improve typing for the 'get' method in session

* update other get signatures. Add tests

Co-authored-by: Federico Caselli <cfederico87@gmail.com>

Author: heckad

sqlalchemy-stubs/ext/asyncio/session.pyi

@@ -8,6 +8,7 @@ from typing import Mapping
 from typing import Optional
 from typing import Protocol
 from typing import Sequence
+from typing import Type
 from typing import TypeVar
 from typing import Union
 
@@ -105,15 +106,15 @@ class _AsyncSessionProtocol(
     ) -> None: ...
     async def get(
         self,
-        entity: Any,
+        entity: Type[_T],
         ident: Any,
         options: Optional[Sequence[Any]] = ...,
         populate_existing: bool = ...,
         with_for_update: Optional[
             Union[Literal[True], Mapping[str, Any]]
         ] = ...,
         identity_token: Optional[Any] = ...,
-    ) -> Any: ...
+    ) -> Optional[_T]: ...
     async def stream(
         self,
         statement: Any,
@@ -155,13 +156,13 @@ class _AsyncSessionTypingCommon(
     async def flush(self, objects: Optional[Any] = ...) -> None: ...
     async def get(
         self,
-        entity: Any,
+        entity: Type[_T],
         ident: Any,
         options: Optional[Any] = ...,
         populate_existing: bool = ...,
         with_for_update: Optional[Any] = ...,
         identity_token: Optional[Any] = ...,
-    ) -> Any: ...
+    ) -> Optional[_T]: ...
     async def merge(self, instance: _T, load: bool = ...) -> _T: ...
     async def refresh(
         self,

sqlalchemy-stubs/orm/session.pyi

@@ -144,15 +144,15 @@ class _SessionProtocol(
     def delete(self, instance: Any) -> None: ...
     def get(
         self,
-        entity: Any,
+        entity: Type[_T],
         ident: Any,
         options: Optional[Sequence[Any]] = ...,
         populate_existing: bool = ...,
         with_for_update: Optional[
             Union[Literal[True], Mapping[str, Any]]
         ] = ...,
         identity_token: Optional[Any] = ...,
-    ) -> Any: ...
+    ) -> Optional[_T]: ...
     def merge(self, instance: _T, load: bool = ...) -> _T: ...
     def flush(self, objects: Optional[Collection[Any]] = ...) -> None: ...
     @classmethod
@@ -364,15 +364,15 @@ class _SessionTypingCommon(_SessionNoIoTypingCommon):
     def flush(self, objects: Optional[Collection[Any]] = ...) -> None: ...
     def get(
         self,
-        entity: Any,
+        entity: Type[_T],
         ident: Any,
         options: Optional[Sequence[Any]] = ...,
         populate_existing: bool = ...,
         with_for_update: Optional[
             Union[Literal[True], Mapping[str, Any]]
         ] = ...,
         identity_token: Optional[Any] = ...,
-    ) -> Any: ...
+    ) -> Optional[_T]: ...
     def bulk_save_objects(
         self,
         objects: Sequence[Any],

test/files/session_get.py

@@ -0,0 +1,24 @@
+from typing import Optional
+
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import registry
+from sqlalchemy.orm import Session
+
+mr: registry = registry()
+
+
+@mr.mapped
+class Foo:
+    id = Column(Integer, primary_key=True)
+    __tablename__ = "foo"
+
+
+s = Session()
+x: Optional[Foo] = s.get(Foo, 1)
+
+
+async def go() -> None:
+    as_ = AsyncSession()
+    y: Optional[Foo] = await as_.get(Foo, 1)