fix(typings): improve typing for server_default

Fixed typing issue where the :paramref:`.AlterColumnOp.server_default` and
:paramref:`.AlterColumnOp.existing_server_default` parameters failed to
accommodate common SQLAlchemy SQL constructs such as ``null()`` and
``text()``.   Pull request courtesy Sebastian Kreft.

this sets up a standard type for the server default argument using
an alias, and adds modifications to write_pyi for extremely basic
ability to render type aliases (with limitations).

Co-authored-by: Mike Bayer <mike_mp@zzzcomputing.com>
Fixes: #1669
Closes: #1670
Pull-request: #1670
Pull-request-sha: e646464

Change-Id: Id25bf7fd706f91aa637adf9b67f0529f1d7d1080

Author: sk-

alembic/ddl/base.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import functools
+from typing import Any
 from typing import Optional
 from typing import TYPE_CHECKING
 from typing import Union
@@ -14,7 +15,10 @@
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.schema import Column
 from sqlalchemy.schema import DDLElement
+from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.elements import quoted_name
+from sqlalchemy.sql.elements import TextClause
+from sqlalchemy.sql.schema import FetchedValue
 
 from ..util.sqla_compat import _columns_for_constraint  # noqa
 from ..util.sqla_compat import _find_columns  # noqa
@@ -23,20 +27,16 @@
 from ..util.sqla_compat import _table_for_constraint  # noqa
 
 if TYPE_CHECKING:
-    from typing import Any
 
     from sqlalchemy import Computed
     from sqlalchemy import Identity
     from sqlalchemy.sql.compiler import Compiled
     from sqlalchemy.sql.compiler import DDLCompiler
-    from sqlalchemy.sql.elements import TextClause
-    from sqlalchemy.sql.functions import Function
-    from sqlalchemy.sql.schema import FetchedValue
     from sqlalchemy.sql.type_api import TypeEngine
 
     from .impl import DefaultImpl
 
-_ServerDefault = Union["TextClause", "FetchedValue", "Function[Any]", str]
+_ServerDefaultType = Union[FetchedValue, str, TextClause, ColumnElement[Any]]
 
 
 class AlterTable(DDLElement):
@@ -75,7 +75,7 @@ def __init__(
         schema: Optional[str] = None,
         existing_type: Optional[TypeEngine] = None,
         existing_nullable: Optional[bool] = None,
-        existing_server_default: Optional[_ServerDefault] = None,
+        existing_server_default: Optional[_ServerDefaultType] = None,
         existing_comment: Optional[str] = None,
     ) -> None:
         super().__init__(name, schema=schema)
@@ -119,7 +119,7 @@ def __init__(
         self,
         name: str,
         column_name: str,
-        default: Optional[_ServerDefault],
+        default: Optional[_ServerDefaultType],
         **kw,
     ) -> None:
         super().__init__(name, column_name, **kw)
@@ -308,7 +308,7 @@ def format_column_name(
 
 def format_server_default(
     compiler: DDLCompiler,
-    default: Optional[_ServerDefault],
+    default: Optional[_ServerDefaultType],
 ) -> str:
     # this can be updated to use compiler.render_default_string
     # for SQLAlchemy 2.0 and above; not in 1.4

alembic/ddl/impl.py

@@ -58,7 +58,7 @@
     from sqlalchemy.sql.selectable import TableClause
     from sqlalchemy.sql.type_api import TypeEngine
 
-    from .base import _ServerDefault
+    from .base import _ServerDefaultType
     from ..autogenerate.api import AutogenContext
     from ..operations.batch import ApplyBatchImpl
     from ..operations.batch import BatchOperationsImpl
@@ -269,7 +269,7 @@ def alter_column(
         *,
         nullable: Optional[bool] = None,
         server_default: Optional[
-            Union[_ServerDefault, Literal[False]]
+            Union[_ServerDefaultType, Literal[False]]
         ] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
@@ -278,7 +278,9 @@ def alter_column(
         comment: Optional[Union[str, Literal[False]]] = False,
         existing_comment: Optional[str] = None,
         existing_type: Optional[TypeEngine] = None,
-        existing_server_default: Optional[_ServerDefault] = None,
+        existing_server_default: Optional[
+            Union[_ServerDefaultType, Literal[False]]
+        ] = None,
         existing_nullable: Optional[bool] = None,
         existing_autoincrement: Optional[bool] = None,
         **kw: Any,

alembic/ddl/mssql.py

@@ -46,7 +46,7 @@
     from sqlalchemy.sql.selectable import TableClause
     from sqlalchemy.sql.type_api import TypeEngine
 
-    from .base import _ServerDefault
+    from .base import _ServerDefaultType
     from .impl import _ReflectedConstraint
 
 
@@ -92,14 +92,14 @@ def alter_column(
         *,
         nullable: Optional[bool] = None,
         server_default: Optional[
-            Union[_ServerDefault, Literal[False]]
+            Union[_ServerDefaultType, Literal[False]]
         ] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
         schema: Optional[str] = None,
         existing_type: Optional[TypeEngine] = None,
         existing_server_default: Union[
-            _ServerDefault, Literal[False], None
+            _ServerDefaultType, Literal[False], None
         ] = None,
         existing_nullable: Optional[bool] = None,
         **kw: Any,

alembic/ddl/mysql.py

@@ -38,7 +38,7 @@
     from sqlalchemy.sql.schema import Constraint
     from sqlalchemy.sql.type_api import TypeEngine
 
-    from .base import _ServerDefault
+    from .base import _ServerDefaultType
 
 
 class MySQLImpl(DefaultImpl):
@@ -83,13 +83,15 @@ def alter_column(
         *,
         nullable: Optional[bool] = None,
         server_default: Optional[
-            Union[_ServerDefault, Literal[False]]
+            Union[_ServerDefaultType, Literal[False]]
         ] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
         schema: Optional[str] = None,
         existing_type: Optional[TypeEngine] = None,
-        existing_server_default: Optional[_ServerDefault] = None,
+        existing_server_default: Optional[
+            Union[_ServerDefaultType, Literal[False]]
+        ] = None,
         existing_nullable: Optional[bool] = None,
         autoincrement: Optional[bool] = None,
         existing_autoincrement: Optional[bool] = None,
@@ -207,7 +209,7 @@ def drop_constraint(
     def _is_mysql_allowed_functional_default(
         self,
         type_: Optional[TypeEngine],
-        server_default: Optional[Union[_ServerDefault, Literal[False]]],
+        server_default: Optional[Union[_ServerDefaultType, Literal[False]]],
     ) -> bool:
         return (
             type_ is not None
@@ -358,7 +360,7 @@ def __init__(
         self,
         name: str,
         column_name: str,
-        default: Optional[_ServerDefault],
+        default: Optional[_ServerDefaultType],
         schema: Optional[str] = None,
     ) -> None:
         super(AlterColumn, self).__init__(name, schema=schema)
@@ -375,7 +377,7 @@ def __init__(
         newname: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
         nullable: Optional[bool] = None,
-        default: Optional[Union[_ServerDefault, Literal[False]]] = False,
+        default: Optional[Union[_ServerDefaultType, Literal[False]]] = False,
         autoincrement: Optional[bool] = None,
         comment: Optional[Union[str, Literal[False]]] = False,
     ) -> None:
@@ -464,7 +466,7 @@ def _mysql_change_column(
 def _mysql_colspec(
     compiler: MySQLDDLCompiler,
     nullable: Optional[bool],
-    server_default: Optional[Union[_ServerDefault, Literal[False]]],
+    server_default: Optional[Union[_ServerDefaultType, Literal[False]]],
     type_: TypeEngine,
     autoincrement: Optional[bool],
     comment: Optional[Union[str, Literal[False]]],

alembic/ddl/postgresql.py

@@ -70,7 +70,7 @@
     from sqlalchemy.sql.schema import Table
     from sqlalchemy.sql.type_api import TypeEngine
 
-    from .base import _ServerDefault
+    from .base import _ServerDefaultType
     from .impl import _ReflectedConstraint
     from ..autogenerate.api import AutogenContext
     from ..autogenerate.render import _f_name
@@ -164,14 +164,16 @@ def alter_column(
         *,
         nullable: Optional[bool] = None,
         server_default: Optional[
-            Union[_ServerDefault, Literal[False]]
+            Union[_ServerDefaultType, Literal[False]]
         ] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
         schema: Optional[str] = None,
         autoincrement: Optional[bool] = None,
         existing_type: Optional[TypeEngine] = None,
-        existing_server_default: Optional[_ServerDefault] = None,
+        existing_server_default: Optional[
+            Union[_ServerDefaultType, Literal[False]]
+        ] = None,
         existing_nullable: Optional[bool] = None,
         existing_autoincrement: Optional[bool] = None,
         **kw: Any,

alembic/op.pyi

@@ -28,13 +28,12 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.expression import TableClause
     from sqlalchemy.sql.schema import Column
-    from sqlalchemy.sql.schema import Computed
-    from sqlalchemy.sql.schema import Identity
     from sqlalchemy.sql.schema import SchemaItem
     from sqlalchemy.sql.schema import Table
     from sqlalchemy.sql.type_api import TypeEngine
     from sqlalchemy.util import immutabledict
 
+    from .ddl.base import _ServerDefaultType
     from .operations.base import BatchOperations
     from .operations.ops import AddColumnOp
     from .operations.ops import AddConstraintOp
@@ -154,14 +153,12 @@ def alter_column(
     *,
     nullable: Optional[bool] = None,
     comment: Union[str, Literal[False], None] = False,
-    server_default: Union[
-        str, bool, Identity, Computed, TextClause, None
-    ] = False,
+    server_default: Union[_ServerDefaultType, None, Literal[False]] = False,
     new_column_name: Optional[str] = None,
     type_: Union[TypeEngine[Any], Type[TypeEngine[Any]], None] = None,
     existing_type: Union[TypeEngine[Any], Type[TypeEngine[Any]], None] = None,
     existing_server_default: Union[
-        str, bool, Identity, Computed, TextClause, None
+        _ServerDefaultType, None, Literal[False]
     ] = False,
     existing_nullable: Optional[bool] = None,
     existing_comment: Optional[str] = None,

alembic/operations/base.py

@@ -27,13 +27,13 @@
 from . import batch
 from . import schemaobj
 from .. import util
+from ..ddl.base import _ServerDefaultType
 from ..util import sqla_compat
 from ..util.compat import formatannotation_fwdref
 from ..util.compat import inspect_formatargspec
 from ..util.compat import inspect_getfullargspec
 from ..util.sqla_compat import _literal_bindparam
 
-
 if TYPE_CHECKING:
     from typing import Literal
 
@@ -44,8 +44,6 @@
     from sqlalchemy.sql.expression import TableClause
     from sqlalchemy.sql.expression import TextClause
     from sqlalchemy.sql.schema import Column
-    from sqlalchemy.sql.schema import Computed
-    from sqlalchemy.sql.schema import Identity
     from sqlalchemy.sql.schema import SchemaItem
     from sqlalchemy.types import TypeEngine
 
@@ -724,15 +722,15 @@ def alter_column(
             nullable: Optional[bool] = None,
             comment: Union[str, Literal[False], None] = False,
             server_default: Union[
-                str, bool, Identity, Computed, TextClause, None
+                _ServerDefaultType, None, Literal[False]
             ] = False,
             new_column_name: Optional[str] = None,
             type_: Union[TypeEngine[Any], Type[TypeEngine[Any]], None] = None,
             existing_type: Union[
                 TypeEngine[Any], Type[TypeEngine[Any]], None
             ] = None,
             existing_server_default: Union[
-                str, bool, Identity, Computed, TextClause, None
+                _ServerDefaultType, None, Literal[False]
             ] = False,
             existing_nullable: Optional[bool] = None,
             existing_comment: Optional[str] = None,
@@ -1691,14 +1689,16 @@ def alter_column(
             *,
             nullable: Optional[bool] = None,
             comment: Union[str, Literal[False], None] = False,
-            server_default: Any = False,
+            server_default: Union[
+                _ServerDefaultType, None, Literal[False]
+            ] = False,
             new_column_name: Optional[str] = None,
             type_: Union[TypeEngine[Any], Type[TypeEngine[Any]], None] = None,
             existing_type: Union[
                 TypeEngine[Any], Type[TypeEngine[Any]], None
             ] = None,
             existing_server_default: Union[
-                str, bool, Identity, Computed, None
+                _ServerDefaultType, None, Literal[False]
             ] = False,
             existing_nullable: Optional[bool] = None,
             existing_comment: Optional[str] = None,

alembic/operations/batch.py

@@ -44,10 +44,10 @@
     from sqlalchemy.engine import Dialect
     from sqlalchemy.sql.elements import ColumnClause
     from sqlalchemy.sql.elements import quoted_name
-    from sqlalchemy.sql.functions import Function
     from sqlalchemy.sql.schema import Constraint
     from sqlalchemy.sql.type_api import TypeEngine
 
+    from ..ddl.base import _ServerDefaultType
     from ..ddl.impl import DefaultImpl
 
 
@@ -485,7 +485,9 @@ def alter_column(
         table_name: str,
         column_name: str,
         nullable: Optional[bool] = None,
-        server_default: Optional[Union[Function[Any], str, bool]] = False,
+        server_default: Union[
+            _ServerDefaultType, None, Literal[False]
+        ] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
         autoincrement: Optional[Union[bool, Literal["auto"]]] = None,

alembic/operations/ops.py

@@ -39,10 +39,8 @@
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.schema import CheckConstraint
     from sqlalchemy.sql.schema import Column
-    from sqlalchemy.sql.schema import Computed
     from sqlalchemy.sql.schema import Constraint
     from sqlalchemy.sql.schema import ForeignKeyConstraint
-    from sqlalchemy.sql.schema import Identity
     from sqlalchemy.sql.schema import Index
     from sqlalchemy.sql.schema import MetaData
     from sqlalchemy.sql.schema import PrimaryKeyConstraint
@@ -53,6 +51,7 @@
     from sqlalchemy.sql.type_api import TypeEngine
 
     from ..autogenerate.rewriter import Rewriter
+    from ..ddl.base import _ServerDefaultType
     from ..runtime.migration import MigrationContext
     from ..script.revision import _RevIdType
 
@@ -1696,7 +1695,9 @@ def __init__(
         *,
         schema: Optional[str] = None,
         existing_type: Optional[Any] = None,
-        existing_server_default: Any = False,
+        existing_server_default: Union[
+            _ServerDefaultType, None, Literal[False]
+        ] = False,
         existing_nullable: Optional[bool] = None,
         existing_comment: Optional[str] = None,
         modify_nullable: Optional[bool] = None,
@@ -1856,15 +1857,15 @@ def alter_column(
         nullable: Optional[bool] = None,
         comment: Optional[Union[str, Literal[False]]] = False,
         server_default: Union[
-            str, bool, Identity, Computed, TextClause, None
+            _ServerDefaultType, None, Literal[False]
         ] = False,
         new_column_name: Optional[str] = None,
         type_: Optional[Union[TypeEngine[Any], Type[TypeEngine[Any]]]] = None,
         existing_type: Optional[
             Union[TypeEngine[Any], Type[TypeEngine[Any]]]
         ] = None,
         existing_server_default: Union[
-            str, bool, Identity, Computed, TextClause, None
+            _ServerDefaultType, None, Literal[False]
         ] = False,
         existing_nullable: Optional[bool] = None,
         existing_comment: Optional[str] = None,
@@ -1980,14 +1981,16 @@ def batch_alter_column(
         *,
         nullable: Optional[bool] = None,
         comment: Optional[Union[str, Literal[False]]] = False,
-        server_default: Any = False,
+        server_default: Union[
+            _ServerDefaultType, None, Literal[False]
+        ] = False,
         new_column_name: Optional[str] = None,
         type_: Optional[Union[TypeEngine[Any], Type[TypeEngine[Any]]]] = None,
         existing_type: Optional[
             Union[TypeEngine[Any], Type[TypeEngine[Any]]]
         ] = None,
-        existing_server_default: Optional[
-            Union[str, bool, Identity, Computed]
+        existing_server_default: Union[
+            _ServerDefaultType, None, Literal[False]
         ] = False,
         existing_nullable: Optional[bool] = None,
         existing_comment: Optional[str] = None,