From e4fa1745839c5a793b2ef2d04d9077f5be65f400 Mon Sep 17 00:00:00 2001 From: Tomasz Nowacki Date: Mon, 4 Mar 2024 09:52:02 -0500 Subject: [PATCH] Fixes: #10933 typing in ColumnExpressionArgument ### Description Fixes: #10933 typing in ColumnExpressionArgument ### Checklist This pull request is: - [ ] A documentation / typographical / small typing error fix - Good to go, no issue or tests are needed - [x] A short code fix - please include the issue number, and create an issue if none exists, which must include a complete example of the issue. one line code fixes without an issue and demonstration will not be accepted. - Please include: `Fixes: #` in the commit message - please include tests. one line code fixes without tests will not be accepted. - [ ] A new feature implementation - please include the issue number, and create an issue if none exists, which must include a complete example of how the feature would look. - Please include: `Fixes: #` in the commit message - please include tests. **Have a nice day!** Closes: #10959 Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/10959 Pull-request-sha: 6fed2cf1d1ba78e9101a3608bd0cf70f2abb3232 Change-Id: I43420add824881e7cc0ec93e3c8b9a04d33e30df --- lib/sqlalchemy/orm/_orm_constructors.py | 5 ++++- lib/sqlalchemy/orm/util.py | 5 ++++- lib/sqlalchemy/sql/lambdas.py | 2 +- test/typing/plain_files/orm/orm_querying.py | 10 ++++++++++ 4 files changed, 19 insertions(+), 3 deletions(-) diff --git a/lib/sqlalchemy/orm/_orm_constructors.py b/lib/sqlalchemy/orm/_orm_constructors.py index f2c4f8ef42..6cf16507ba 100644 --- a/lib/sqlalchemy/orm/_orm_constructors.py +++ b/lib/sqlalchemy/orm/_orm_constructors.py @@ -722,7 +722,10 @@ def composite( def with_loader_criteria( entity_or_base: _EntityType[Any], - where_criteria: _ColumnExpressionArgument[bool], + where_criteria: Union[ + _ColumnExpressionArgument[bool], + Callable[[Any], _ColumnExpressionArgument[bool]], + ], loader_only: bool = False, include_aliases: bool = False, propagate_to_loaders: bool = True, diff --git a/lib/sqlalchemy/orm/util.py b/lib/sqlalchemy/orm/util.py index 1fd0f6863d..297e556ab3 100644 --- a/lib/sqlalchemy/orm/util.py +++ b/lib/sqlalchemy/orm/util.py @@ -1380,7 +1380,10 @@ class LoaderCriteriaOption(CriteriaOption): def __init__( self, entity_or_base: _EntityType[Any], - where_criteria: _ColumnExpressionArgument[bool], + where_criteria: Union[ + _ColumnExpressionArgument[bool], + Callable[[Any], _ColumnExpressionArgument[bool]], + ], loader_only: bool = False, include_aliases: bool = False, propagate_to_loaders: bool = True, diff --git a/lib/sqlalchemy/sql/lambdas.py b/lib/sqlalchemy/sql/lambdas.py index 726fa2411f..7a6b7b8f77 100644 --- a/lib/sqlalchemy/sql/lambdas.py +++ b/lib/sqlalchemy/sql/lambdas.py @@ -437,7 +437,7 @@ class DeferredLambdaElement(LambdaElement): def __init__( self, - fn: _LambdaType, + fn: _AnyLambdaType, role: Type[roles.SQLRole], opts: Union[Type[LambdaOptions], LambdaOptions] = LambdaOptions, lambda_args: Tuple[Any, ...] = (), diff --git a/test/typing/plain_files/orm/orm_querying.py b/test/typing/plain_files/orm/orm_querying.py index fa59baad43..3251147dd6 100644 --- a/test/typing/plain_files/orm/orm_querying.py +++ b/test/typing/plain_files/orm/orm_querying.py @@ -1,5 +1,6 @@ from __future__ import annotations +from sqlalchemy import ColumnElement from sqlalchemy import ForeignKey from sqlalchemy import orm from sqlalchemy import select @@ -124,3 +125,12 @@ def load_options_error() -> None: # EXPECTED_MYPY_RE: Argument 1 to .* has incompatible type .* orm.undefer(B.a).undefer("bar"), ) + + +# test 10959 +def test_10959_with_loader_criteria() -> None: + def where_criteria(cls_: type[A]) -> ColumnElement[bool]: + return cls_.data == "some data" + + orm.with_loader_criteria(A, lambda cls: cls.data == "some data") + orm.with_loader_criteria(A, where_criteria) -- 2.47.2