Merge "Correct type hint for FunctionElement.table_valued()" into main

whysage · Jan 24, 2024 · ae2aefb · ae2aefb
2 parents abadb14 + 48d3ad2
commit ae2aefb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
@@ -68,6 +68,7 @@
     from ._typing import _ByArgument
     from ._typing import _ColumnExpressionArgument
     from ._typing import _ColumnExpressionOrLiteralArgument
+    from ._typing import _ColumnExpressionOrStrLabelArgument
     from ._typing import _TypeEngineArgument
     from .base import _EntityNamespace
     from .elements import ClauseElement
@@ -235,7 +236,7 @@ def scalar_table_valued(
         return ScalarFunctionColumn(self, name, type_)
 
     def table_valued(
-        self, *expr: _ColumnExpressionArgument[Any], **kw: Any
+        self, *expr: _ColumnExpressionOrStrLabelArgument[Any], **kw: Any
     ) -> TableValuedAlias:
         r"""Return a :class:`_sql.TableValuedAlias` representation of this
         :class:`_functions.FunctionElement` with table-valued expressions added.

diff --git a/test/typing/plain_files/sql/functions_again.py b/test/typing/plain_files/sql/functions_again.py
@@ -54,3 +54,9 @@ class Foo(Base):
 ).group_by(Foo.a)
 # EXPECTED_TYPE: Select[Tuple[int, str]]
 reveal_type(stmt2)
+
+
+# EXPECTED_TYPE: TableValuedAlias
+reveal_type(func.json_each().table_valued("key", "value"))
+# EXPECTED_TYPE: TableValuedAlias
+reveal_type(func.json_each().table_valued(Foo.a, Foo.b))