From ceb0eb44ce2da4b7161b7a115525ed6bccf08cbc Mon Sep 17 00:00:00 2001 From: Federico Caselli Date: Tue, 31 Jan 2023 21:20:26 +0100 Subject: [PATCH] Fixed typing of limit, offset and fetch to allow ``None``. Fixes: #9183 Change-Id: I1ac3e3698034826122ea8a0cdc9f8f55a10ed6c1 --- doc/build/changelog/unreleased_20/9183.rst | 7 +++++++ lib/sqlalchemy/orm/query.py | 9 +++------ lib/sqlalchemy/sql/_typing.py | 2 ++ lib/sqlalchemy/sql/selectable.py | 11 +++++------ lib/sqlalchemy/sql/util.py | 12 ++++++------ test/ext/mypy/plain_files/session.py | 4 ++++ test/ext/mypy/plain_files/typed_queries.py | 14 +++++++++++++- 7 files changed, 40 insertions(+), 19 deletions(-) create mode 100644 doc/build/changelog/unreleased_20/9183.rst diff --git a/doc/build/changelog/unreleased_20/9183.rst b/doc/build/changelog/unreleased_20/9183.rst new file mode 100644 index 0000000000..76e306356c --- /dev/null +++ b/doc/build/changelog/unreleased_20/9183.rst @@ -0,0 +1,7 @@ +.. change:: + :tags: bug, typing + :tickets: 9183 + + Fixed typing of limit, offset and fetch to allow ``None``. + + diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py index ad4b3abcf5..4c182ebe66 100644 --- a/lib/sqlalchemy/orm/query.py +++ b/lib/sqlalchemy/orm/query.py @@ -116,6 +116,7 @@ if TYPE_CHECKING: from ..sql._typing import _ColumnsClauseArgument from ..sql._typing import _DMLColumnArgument from ..sql._typing import _JoinTargetArgument + from ..sql._typing import _LimitOffsetType from ..sql._typing import _MAYBE_ENTITY from ..sql._typing import _no_kw from ..sql._typing import _NOT_ENTITY @@ -2615,9 +2616,7 @@ class Query( @_generative @_assertions(_no_statement_condition) - def limit( - self: SelfQuery, limit: Union[int, _ColumnExpressionArgument[int]] - ) -> SelfQuery: + def limit(self: SelfQuery, limit: _LimitOffsetType) -> SelfQuery: """Apply a ``LIMIT`` to the query and return the newly resulting ``Query``. @@ -2631,9 +2630,7 @@ class Query( @_generative @_assertions(_no_statement_condition) - def offset( - self: SelfQuery, offset: Union[int, _ColumnExpressionArgument[int]] - ) -> SelfQuery: + def offset(self: SelfQuery, offset: _LimitOffsetType) -> SelfQuery: """Apply an ``OFFSET`` to the query and return the newly resulting ``Query``. diff --git a/lib/sqlalchemy/sql/_typing.py b/lib/sqlalchemy/sql/_typing.py index 5c1a501e47..e1190f7dd2 100644 --- a/lib/sqlalchemy/sql/_typing.py +++ b/lib/sqlalchemy/sql/_typing.py @@ -260,6 +260,8 @@ _TypeEngineArgument = Union[Type["TypeEngine[_T]"], "TypeEngine[_T]"] _EquivalentColumnMap = Dict["ColumnElement[Any]", Set["ColumnElement[Any]"]] +_LimitOffsetType = Union[int, _ColumnExpressionArgument[int], None] + if TYPE_CHECKING: def is_sql_compiler(c: Compiled) -> TypeGuard[SQLCompiler]: diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py index f43e6b43f8..47cf683577 100644 --- a/lib/sqlalchemy/sql/selectable.py +++ b/lib/sqlalchemy/sql/selectable.py @@ -109,6 +109,7 @@ if TYPE_CHECKING: from ._typing import _ColumnExpressionOrStrLabelArgument from ._typing import _FromClauseArgument from ._typing import _JoinTargetArgument + from ._typing import _LimitOffsetType from ._typing import _MAYBE_ENTITY from ._typing import _NOT_ENTITY from ._typing import _OnClauseArgument @@ -3955,7 +3956,7 @@ class GenerativeSelect(SelectBase, Generative): def _offset_or_limit_clause( self, - element: Union[int, _ColumnExpressionArgument[Any]], + element: _LimitOffsetType, name: Optional[str] = None, type_: Optional[_TypeEngineArgument[int]] = None, ) -> ColumnElement[Any]: @@ -4041,8 +4042,7 @@ class GenerativeSelect(SelectBase, Generative): @_generative def limit( - self: SelfGenerativeSelect, - limit: Union[int, _ColumnExpressionArgument[int]], + self: SelfGenerativeSelect, limit: _LimitOffsetType ) -> SelfGenerativeSelect: """Return a new selectable with the given LIMIT criterion applied. @@ -4078,7 +4078,7 @@ class GenerativeSelect(SelectBase, Generative): @_generative def fetch( self: SelfGenerativeSelect, - count: Union[int, _ColumnExpressionArgument[int]], + count: _LimitOffsetType, with_ties: bool = False, percent: bool = False, ) -> SelfGenerativeSelect: @@ -4133,8 +4133,7 @@ class GenerativeSelect(SelectBase, Generative): @_generative def offset( - self: SelfGenerativeSelect, - offset: Union[int, _ColumnExpressionArgument[int]], + self: SelfGenerativeSelect, offset: _LimitOffsetType ) -> SelfGenerativeSelect: """Return a new selectable with the given OFFSET criterion applied. diff --git a/lib/sqlalchemy/sql/util.py b/lib/sqlalchemy/sql/util.py index a92ee9d1a0..1dad9ce684 100644 --- a/lib/sqlalchemy/sql/util.py +++ b/lib/sqlalchemy/sql/util.py @@ -71,8 +71,8 @@ from ..util.typing import Literal from ..util.typing import Protocol if typing.TYPE_CHECKING: - from ._typing import _ColumnExpressionArgument from ._typing import _EquivalentColumnMap + from ._typing import _LimitOffsetType from ._typing import _TypeEngineArgument from .elements import BinaryExpression from .elements import TextClause @@ -1411,7 +1411,7 @@ class ColumnAdapter(ClauseAdapter): def _offset_or_limit_clause( - element: Union[int, _ColumnExpressionArgument[int]], + element: _LimitOffsetType, name: Optional[str] = None, type_: Optional[_TypeEngineArgument[int]] = None, ) -> ColumnElement[int]: @@ -1427,8 +1427,8 @@ def _offset_or_limit_clause( def _offset_or_limit_clause_asint_if_possible( - clause: Optional[Union[int, _ColumnExpressionArgument[int]]] -) -> Optional[Union[int, _ColumnExpressionArgument[int]]]: + clause: _LimitOffsetType, +) -> _LimitOffsetType: """Return the offset or limit clause as a simple integer if possible, else return the clause. @@ -1443,8 +1443,8 @@ def _offset_or_limit_clause_asint_if_possible( def _make_slice( - limit_clause: Optional[Union[int, _ColumnExpressionArgument[int]]], - offset_clause: Optional[Union[int, _ColumnExpressionArgument[int]]], + limit_clause: _LimitOffsetType, + offset_clause: _LimitOffsetType, start: int, stop: int, ) -> Tuple[Optional[ColumnElement[int]], Optional[ColumnElement[int]]]: diff --git a/test/ext/mypy/plain_files/session.py b/test/ext/mypy/plain_files/session.py index 636e3854a5..9106b90169 100644 --- a/test/ext/mypy/plain_files/session.py +++ b/test/ext/mypy/plain_files/session.py @@ -89,4 +89,8 @@ with Session(e) as sess: # EXPECTED_TYPE: User reveal_type(uobj1) + sess.query(User).limit(None).offset(None).limit(10).offset(10).limit( + User.id + ).offset(User.id) + # more result tests in typed_results.py diff --git a/test/ext/mypy/plain_files/typed_queries.py b/test/ext/mypy/plain_files/typed_queries.py index fb988c9851..3e67a71325 100644 --- a/test/ext/mypy/plain_files/typed_queries.py +++ b/test/ext/mypy/plain_files/typed_queries.py @@ -48,7 +48,19 @@ def t_select_1() -> None: def t_select_2() -> None: - stmt = select(User).filter(User.id == 5) + stmt = ( + select(User) + .filter(User.id == 5) + .limit(1) + .offset(3) + .offset(None) + .limit(None) + .limit(User.id) + .offset(User.id) + .fetch(1) + .fetch(None) + .fetch(User.id) + ) # EXPECTED_TYPE: Select[Tuple[User]] reveal_type(stmt) -- 2.47.2