Fixed typing of limit, offset and fetch to allow ``None``.

Fixes: #9183
Change-Id: I1ac3e3698034826122ea8a0cdc9f8f55a10ed6c1

diff --git a/doc/build/changelog/unreleased_20/9183.rst b/doc/build/changelog/unreleased_20/9183.rst
new file mode 100644 (file)
index 0000000..76e3063
--- /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 ad4b3abcf56ea21ca7b7d71a1ed535e803e66526..4c182ebe668d496ae180b015a9b787409ba3f26c 100644 (file)
--- 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 5c1a501e47e0bf3b29b030360a91393704fa29f8..e1190f7dd294c9fc760e6c5968a1e21fa39e5a13 100644 (file)
--- 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 f43e6b43f82bd8daf168451399d2dd257c85b8fd..47cf68357782cb849b52322a2a75517079eb044d 100644 (file)
--- 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 a92ee9d1a06ec22cfaeb7429c6f738b3f6054c10..1dad9ce6846d478c9774eb4aed93f94ab2453390 100644 (file)
--- 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 636e3854a51da3cdd0255b422df578d11c6833d7..9106b901690944ec81b80dc7771c110a3462c4a2 100644 (file)
--- 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 fb988c985194b19b0c0c14e86b11faaf62669260..3e67a71325b91dae691a0682712d1467a32be979 100644 (file)
--- 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)