🎨 Fix type detection of select results in PyCharm (#15)

diff --git a/sqlmodel/orm/session.py b/sqlmodel/orm/session.py
index d708f1e7e9dc80a70684f12bc599be183394c20e..a5a63e2c69b2e342967ea1f2dfa8627b0e6d48db 100644 (file)
--- a/sqlmodel/orm/session.py
+++ b/sqlmodel/orm/session.py
@@ -10,14 +10,14 @@ from typing_extensions import Literal
 from ..engine.result import Result, ScalarResult
 from ..sql.base import Executable
 
-_T = TypeVar("_T")
+_TSelectParam = TypeVar("_TSelectParam")
 
 
 class Session(_Session):
     @overload
     def exec(
         self,
-        statement: Select[_T],
+        statement: Select[_TSelectParam],
         *,
         params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
         execution_options: Mapping[str, Any] = util.EMPTY_DICT,
@@ -25,13 +25,13 @@ class Session(_Session):
         _parent_execute_state: Optional[Any] = None,
         _add_event: Optional[Any] = None,
         **kw: Any,
-    ) -> Union[Result[_T]]:
+    ) -> Result[_TSelectParam]:
         ...
 
     @overload
     def exec(
         self,
-        statement: SelectOfScalar[_T],
+        statement: SelectOfScalar[_TSelectParam],
         *,
         params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
         execution_options: Mapping[str, Any] = util.EMPTY_DICT,
@@ -39,12 +39,16 @@ class Session(_Session):
         _parent_execute_state: Optional[Any] = None,
         _add_event: Optional[Any] = None,
         **kw: Any,
-    ) -> Union[ScalarResult[_T]]:
+    ) -> ScalarResult[_TSelectParam]:
         ...
 
     def exec(
         self,
-        statement: Union[Select[_T], SelectOfScalar[_T], Executable[_T]],
+        statement: Union[
+            Select[_TSelectParam],
+            SelectOfScalar[_TSelectParam],
+            Executable[_TSelectParam],
+        ],
         *,
         params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
         execution_options: Mapping[str, Any] = util.EMPTY_DICT,
@@ -52,7 +56,7 @@ class Session(_Session):
         _parent_execute_state: Optional[Any] = None,
         _add_event: Optional[Any] = None,
         **kw: Any,
-    ) -> Union[Result[_T], ScalarResult[_T]]:
+    ) -> Union[Result[_TSelectParam], ScalarResult[_TSelectParam]]:
         results = super().execute(
             statement,
             params=params,
@@ -118,13 +122,13 @@ class Session(_Session):
 
     def get(
         self,
-        entity: Type[_T],
+        entity: Type[_TSelectParam],
         ident: Any,
         options: Optional[Sequence[Any]] = None,
         populate_existing: bool = False,
         with_for_update: Optional[Union[Literal[True], Mapping[str, Any]]] = None,
         identity_token: Optional[Any] = None,
-    ) -> Optional[_T]:
+    ) -> Optional[_TSelectParam]:
         return super().get(
             entity,
             ident,