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

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

View File

@ -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,