Skip to content

Commit 1da849a

Browse files
authored
🎨 Fix type detection of select results in PyCharm (#15)
1 parent af03df8 commit 1da849a

File tree

1 file changed

+13
-9
lines changed

1 file changed

+13
-9
lines changed

sqlmodel/orm/session.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,49 +10,53 @@
1010
from ..engine.result import Result, ScalarResult
1111
from ..sql.base import Executable
1212

13-
_T = TypeVar("_T")
13+
_TSelectParam = TypeVar("_TSelectParam")
1414

1515

1616
class Session(_Session):
1717
@overload
1818
def exec(
1919
self,
20-
statement: Select[_T],
20+
statement: Select[_TSelectParam],
2121
*,
2222
params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
2323
execution_options: Mapping[str, Any] = util.EMPTY_DICT,
2424
bind_arguments: Optional[Mapping[str, Any]] = None,
2525
_parent_execute_state: Optional[Any] = None,
2626
_add_event: Optional[Any] = None,
2727
**kw: Any,
28-
) -> Union[Result[_T]]:
28+
) -> Result[_TSelectParam]:
2929
...
3030

3131
@overload
3232
def exec(
3333
self,
34-
statement: SelectOfScalar[_T],
34+
statement: SelectOfScalar[_TSelectParam],
3535
*,
3636
params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
3737
execution_options: Mapping[str, Any] = util.EMPTY_DICT,
3838
bind_arguments: Optional[Mapping[str, Any]] = None,
3939
_parent_execute_state: Optional[Any] = None,
4040
_add_event: Optional[Any] = None,
4141
**kw: Any,
42-
) -> Union[ScalarResult[_T]]:
42+
) -> ScalarResult[_TSelectParam]:
4343
...
4444

4545
def exec(
4646
self,
47-
statement: Union[Select[_T], SelectOfScalar[_T], Executable[_T]],
47+
statement: Union[
48+
Select[_TSelectParam],
49+
SelectOfScalar[_TSelectParam],
50+
Executable[_TSelectParam],
51+
],
4852
*,
4953
params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
5054
execution_options: Mapping[str, Any] = util.EMPTY_DICT,
5155
bind_arguments: Optional[Mapping[str, Any]] = None,
5256
_parent_execute_state: Optional[Any] = None,
5357
_add_event: Optional[Any] = None,
5458
**kw: Any,
55-
) -> Union[Result[_T], ScalarResult[_T]]:
59+
) -> Union[Result[_TSelectParam], ScalarResult[_TSelectParam]]:
5660
results = super().execute(
5761
statement,
5862
params=params,
@@ -118,13 +122,13 @@ def query(self, *entities: Any, **kwargs: Any) -> "_Query[Any]":
118122

119123
def get(
120124
self,
121-
entity: Type[_T],
125+
entity: Type[_TSelectParam],
122126
ident: Any,
123127
options: Optional[Sequence[Any]] = None,
124128
populate_existing: bool = False,
125129
with_for_update: Optional[Union[Literal[True], Mapping[str, Any]]] = None,
126130
identity_token: Optional[Any] = None,
127-
) -> Optional[_T]:
131+
) -> Optional[_TSelectParam]:
128132
return super().get(
129133
entity,
130134
ident,

0 commit comments

Comments
 (0)