improve types of over in functions (#158)

Author: CaselIT

sqlalchemy-stubs/sql/elements.pyi

@@ -59,6 +59,8 @@ _BE = TypeVar("_BE", bound=BinaryExpression[Any])
 _QN = TypeVar("_QN", bound=quoted_name)
 _TL = TypeVar("_TL", bound=_truncated_label)
 
+_OverByType = Union[ClauseElement, str]
+
 def collate(expression: Any, collation: str) -> BinaryExpression[_TE]: ...
 def between(
     expr: Any,
@@ -542,8 +544,10 @@ class Over(ColumnElement[_TE]):
     def __init__(
         self,
         element: ColumnElement[_TE],
-        partition_by: Optional[ClauseElement] = ...,
-        order_by: Optional[ClauseElement] = ...,
+        partition_by: Optional[
+            Union[_OverByType, Sequence[_OverByType]]
+        ] = ...,
+        order_by: Optional[Union[_OverByType, Sequence[_OverByType]]] = ...,
         range_: Optional[Any] = ...,
         rows: Optional[Any] = ...,
     ) -> None: ...

sqlalchemy-stubs/sql/functions.pyi

@@ -2,6 +2,7 @@ from typing import Any
 from typing import Generic
 from typing import Optional
 from typing import overload
+from typing import Sequence
 from typing import Type
 from typing import TypeVar
 from typing import Union
@@ -35,6 +36,8 @@ _T_co = TypeVar("_T_co", covariant=True)
 _TE = TypeVar("_TE", bound=type_api.TypeEngine[Any])
 _FE = TypeVar("_FE", bound=FunctionElement[Any])
 
+_OverByType = Union[ClauseElement, str]
+
 def register_function(
     identifier: str, fn: Any, package: str = ...
 ) -> None: ...
@@ -63,8 +66,10 @@ class FunctionElement(  # type: ignore[misc]
     def clauses(self) -> ClauseList[Any]: ...
     def over(
         self,
-        partition_by: Optional[ClauseElement] = ...,
-        order_by: Optional[ClauseElement] = ...,
+        partition_by: Optional[
+            Union[_OverByType, Sequence[_OverByType]]
+        ] = ...,
+        order_by: Optional[Union[_OverByType, Sequence[_OverByType]]] = ...,
         rows: Optional[Any] = ...,
         range_: Optional[Any] = ...,
     ) -> Over[_TE]: ...

test/files/functions.py

@@ -0,0 +1,19 @@
+from sqlalchemy import Column, Integer, create_engine, func
+from sqlalchemy.orm import Session, declarative_base, sessionmaker
+
+Base = declarative_base()
+
+
+class Foo(Base):
+    __tablename__ = "foo"
+
+    id = Column(Integer(), primary_key=True)
+    a = Column(Integer())
+    b = Column(Integer())
+
+
+func.row_number().over(order_by=Foo.a, partition_by=Foo.b.desc())
+func.row_number().over(order_by=[Foo.a.desc(), Foo.b.desc()])
+func.row_number().over(partition_by=[Foo.a.desc(), Foo.b.desc()])
+func.row_number().over(order_by="a", partition_by=("a", "b"))
+func.row_number().over(partition_by="a", order_by=("a", "b"))