Fixes: #10933 typing in ColumnExpressionArgument

### Description
Fixes: #10933 typing in ColumnExpressionArgument
### Checklist
<!-- go over following points. check them with an `x` if they do apply, (they turn into clickable checkboxes once the PR is submitted, so no need to do everything at once)

-->

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: #<issue number>` 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: #<issue number>` 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

diff --git a/lib/sqlalchemy/orm/_orm_constructors.py b/lib/sqlalchemy/orm/_orm_constructors.py
index f2c4f8ef4236a0df6e2c09f2ae5ffb709c5562f8..6cf16507ba6f2686b1072bae9b6d20d8f7867a30 100644 (file)
--- 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 1fd0f6863df29a697056775095bc0c59789f289c..297e556ab3d96565b650e978a4b8b1d1de2afcbc 100644 (file)
--- 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 726fa2411f86b779badec9d0c1073ecb7de54df2..7a6b7b8f776a34e2448c27d6c9dbd2df16563e5f 100644 (file)
--- 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 fa59baad43a9185d3db3aea550834f0ac3801d79..3251147dd6876e04f4263762af2902512a5b4991 100644 (file)
--- 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)