GitHub-Proxy/sqlmodel-fix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sebastián Ramírez 2021-08-25 15:43:53 +02:00 committed by GitHub
parent af03df88ac
commit 1da849ac48
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
sqlmodel/orm/session.py
View File

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