Allow FunctionElement for Column server_onupdate (#152)

Use common type aliases for Column server_default and server_onupdate

Author: jwilges

sqlalchemy-stubs/sql/schema.pyi

@@ -29,6 +29,7 @@ from .elements import ClauseElement
 from .elements import ColumnClause
 from .elements import ColumnElement
 from .elements import TextClause
+from .functions import FunctionElement
 from .events import DDLEvents
 from .selectable import TableClause
 from .. import util
@@ -54,6 +55,9 @@ _IDX = TypeVar("_IDX", bound=Index)
 _CP = TypeVar("_CP", bound=Computed)
 _ID = TypeVar("_ID", bound=Identity)
 
+_ServerDefaultType = Union[FetchedValue, str, TextClause, ColumnElement[_T]]
+_ServerOnUpdateType = Union[FetchedValue, FunctionElement]
+
 class SchemaItem(SchemaEventTarget, visitors.Visitable):
     __visit_name__: str = ...
     create_drop_stringify_dialect: str = ...
@@ -140,10 +144,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
     primary_key: bool = ...
     nullable: bool = ...
     default: Optional[Any] = ...
-    server_default: Optional[
-        Union[FetchedValue, str, TextClause, ColumnElement[_TE]]
-    ] = ...
-    server_onupdate: Optional[FetchedValue] = ...
+    server_default: Optional[_ServerDefaultType[_TE]] = ...
+    server_onupdate: Optional[_ServerOnUpdateType] = ...
     index: Optional[bool] = ...
     unique: Optional[bool] = ...
     system: bool = ...
@@ -169,10 +171,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
         nullable: bool = ...,
         onupdate: Optional[Any] = ...,
         primary_key: bool = ...,
-        server_default: Optional[
-            Union[FetchedValue, str, TextClause, ColumnElement[Any]]
-        ] = ...,
-        server_onupdate: Optional[FetchedValue] = ...,
+        server_default: Optional[_ServerDefaultType[Any]] = ...,
+        server_onupdate: Optional[_ServerOnUpdateType] = ...,
         quote: Optional[bool] = ...,
         unique: Optional[bool] = ...,
         system: bool = ...,
@@ -192,10 +192,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
         nullable: bool = ...,
         onupdate: Optional[Any] = ...,
         primary_key: bool = ...,
-        server_default: Optional[
-            Union[FetchedValue, str, TextClause, ColumnElement[Any]]
-        ] = ...,
-        server_onupdate: Optional[FetchedValue] = ...,
+        server_default: Optional[_ServerDefaultType[Any]] = ...,
+        server_onupdate: Optional[_ServerOnUpdateType] = ...,
         quote: Optional[bool] = ...,
         unique: Optional[bool] = ...,
         system: bool = ...,
@@ -217,10 +215,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
         nullable: bool = ...,
         onupdate: Optional[Any] = ...,
         primary_key: bool = ...,
-        server_default: Optional[
-            Union[FetchedValue, str, TextClause, ColumnElement[_TE]]
-        ] = ...,
-        server_onupdate: Optional[FetchedValue] = ...,
+        server_default: Optional[_ServerDefaultType[_TE]] = ...,
+        server_onupdate: Optional[_ServerOnUpdateType] = ...,
         quote: Optional[bool] = ...,
         unique: Optional[bool] = ...,
         system: bool = ...,
@@ -241,10 +237,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
         nullable: bool = ...,
         onupdate: Optional[Any] = ...,
         primary_key: bool = ...,
-        server_default: Optional[
-            Union[FetchedValue, str, TextClause, ColumnElement[_TE]]
-        ] = ...,
-        server_onupdate: Optional[FetchedValue] = ...,
+        server_default: Optional[_ServerDefaultType[_TE]] = ...,
+        server_onupdate: Optional[_ServerOnUpdateType] = ...,
         quote: Optional[bool] = ...,
         unique: Optional[bool] = ...,
         system: bool = ...,

test/files/column_server_onupdate.py

@@ -0,0 +1,20 @@
+from sqlalchemy import Column
+from sqlalchemy.orm import registry
+from sqlalchemy.sql import functions as func
+from sqlalchemy.types import Integer
+
+reg: registry = registry()
+
+
+@reg.mapped
+class ColumnServerOnUpdateArgument:
+    __tablename__ = "onupdate"
+
+    # overload 1: (__name: str, *args, ...)
+    overload1: int = Column("name", server_onupdate=func.now(), nullable=False)
+    # overload 2: (*args, ...)
+    Column(server_onupdate=func.now(), nullable=False)
+    # overload 3: (__name: str, __type: _TE, *args, ...)
+    overload3 = Column("name", Integer, server_onupdate=func.now(), nullable=False)
+    # overload 4: (__type: _TE, *args, ...)
+    overload4 = Column(Integer, server_onupdate=func.now(), nullable=False)