implement basic typing for lambda elements

These weren't working at all, so fixed things up and
added a test suite.   Keeping things very basic with Any
returns etc. as having more specific return types
starts making it too cumbersome to write end-user code.

Corrected the type passed for "lambda statements" so that a plain lambda is
accepted by mypy, pyright, others without any errors about argument types.
Additionally implemented typing for more of the public API for lambda
statements and ensured :class:`.StatementLambdaElement` is part of the
:class:`.Executable` hierarchy so it's typed as accepted by
:meth:`_engine.Connection.execute`.

Fixes: #9120
Change-Id: Ia7fa34e5b6e43fba02c8f94ccc256f3a68a1f445

1 files changed, 77 insertions, 0 deletions

diff --git a/test/ext/mypy/plain_files/lambda_stmt.py b/test/ext/mypy/plain_files/lambda_stmt.py
new file mode 100644
index 000000000..7e15778c1
--- /dev/null
+++ b/test/ext/mypy/plain_files/lambda_stmt.py
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import Tuple
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column
+from sqlalchemy import create_engine
+from sqlalchemy import Integer
+from sqlalchemy import lambda_stmt
+from sqlalchemy import MetaData
+from sqlalchemy import Result
+from sqlalchemy import select
+from sqlalchemy import String
+from sqlalchemy import Table
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class User(Base):
+    __tablename__ = "a"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    email: Mapped[str]
+
+
+user_table = Table(
+    "user_table", MetaData(), Column("id", Integer), Column("email", String)
+)
+
+
+s1 = select(user_table).where(lambda: user_table.c.id == 5)
+
+s2 = select(User).where(lambda: User.id == 5)
+
+s3 = lambda_stmt(lambda: select(user_table).where(user_table.c.id == 5))
+
+s4 = lambda_stmt(lambda: select(User).where(User.id == 5))
+
+s5 = lambda_stmt(lambda: select(user_table)) + (
+    lambda s: s.where(user_table.c.id == 5)
+)
+
+s6 = lambda_stmt(lambda: select(User)) + (lambda s: s.where(User.id == 5))
+
+
+if TYPE_CHECKING:
+
+    # EXPECTED_TYPE: StatementLambdaElement
+    reveal_type(s5)
+
+    # EXPECTED_TYPE: StatementLambdaElement
+    reveal_type(s6)
+
+
+e = create_engine("sqlite://")
+
+with e.connect() as conn:
+    result = conn.execute(s6)
+
+    if TYPE_CHECKING:
+        # EXPECTED_TYPE: CursorResult[Any]
+        reveal_type(result)
+
+    # we can type these like this
+    my_result: Result[Tuple[User]] = conn.execute(s6)
+
+    if TYPE_CHECKING:
+        # pyright and mypy disagree on the specific type here,
+        # mypy sees Result as we said, pyright seems to upgrade it to
+        # CursorResult
+        # EXPECTED_RE_TYPE: .*(?:Cursor)?Result\[Tuple\[.*User\]\]
+        reveal_type(my_result)