API refactor

venv/lib/python3.12/site-packages/alembic/operations/base.py

@@ -1,3 +1,5 @@
+# mypy: allow-untyped-calls
+
 from __future__ import annotations
 
 from contextlib import contextmanager
@@ -10,7 +12,9 @@ from typing import Dict
 from typing import Iterator
 from typing import List  # noqa
 from typing import Mapping
+from typing import NoReturn
 from typing import Optional
+from typing import overload
 from typing import Sequence  # noqa
 from typing import Tuple
 from typing import Type  # noqa
@@ -39,7 +43,6 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.expression import ColumnElement
     from sqlalchemy.sql.expression import TableClause
     from sqlalchemy.sql.expression import TextClause
-    from sqlalchemy.sql.functions import Function
     from sqlalchemy.sql.schema import Column
     from sqlalchemy.sql.schema import Computed
     from sqlalchemy.sql.schema import Identity
@@ -47,12 +50,28 @@ if TYPE_CHECKING:
     from sqlalchemy.types import TypeEngine
 
     from .batch import BatchOperationsImpl
+    from .ops import AddColumnOp
+    from .ops import AddConstraintOp
+    from .ops import AlterColumnOp
+    from .ops import AlterTableOp
+    from .ops import BulkInsertOp
+    from .ops import CreateIndexOp
+    from .ops import CreateTableCommentOp
+    from .ops import CreateTableOp
+    from .ops import DropColumnOp
+    from .ops import DropConstraintOp
+    from .ops import DropIndexOp
+    from .ops import DropTableCommentOp
+    from .ops import DropTableOp
+    from .ops import ExecuteSQLOp
     from .ops import MigrateOperation
     from ..ddl import DefaultImpl
     from ..runtime.migration import MigrationContext
 __all__ = ("Operations", "BatchOperations")
 _T = TypeVar("_T")
 
+_C = TypeVar("_C", bound=Callable[..., Any])
+
 
 class AbstractOperations(util.ModuleClsProxy):
     """Base class for Operations and BatchOperations.
@@ -86,7 +105,7 @@ class AbstractOperations(util.ModuleClsProxy):
     @classmethod
     def register_operation(
         cls, name: str, sourcename: Optional[str] = None
-    ) -> Callable[[_T], _T]:
+    ) -> Callable[[Type[_T]], Type[_T]]:
         """Register a new operation for this class.
 
         This method is normally used to add new operations
@@ -103,7 +122,7 @@ class AbstractOperations(util.ModuleClsProxy):
 
         """
 
-        def register(op_cls):
+        def register(op_cls: Type[_T]) -> Type[_T]:
             if sourcename is None:
                 fn = getattr(op_cls, name)
                 source_name = fn.__name__
@@ -122,8 +141,11 @@ class AbstractOperations(util.ModuleClsProxy):
                 *spec, formatannotation=formatannotation_fwdref
             )
             num_defaults = len(spec[3]) if spec[3] else 0
+
+            defaulted_vals: Tuple[Any, ...]
+
             if num_defaults:
-                defaulted_vals = name_args[0 - num_defaults :]
+                defaulted_vals = tuple(name_args[0 - num_defaults :])
             else:
                 defaulted_vals = ()
 
@@ -164,7 +186,7 @@ class AbstractOperations(util.ModuleClsProxy):
 
             globals_ = dict(globals())
             globals_.update({"op_cls": op_cls})
-            lcl = {}
+            lcl: Dict[str, Any] = {}
 
             exec(func_text, globals_, lcl)
             setattr(cls, name, lcl[name])
@@ -180,7 +202,7 @@ class AbstractOperations(util.ModuleClsProxy):
         return register
 
     @classmethod
-    def implementation_for(cls, op_cls: Any) -> Callable[..., Any]:
+    def implementation_for(cls, op_cls: Any) -> Callable[[_C], _C]:
         """Register an implementation for a given :class:`.MigrateOperation`.
 
         This is part of the operation extensibility API.
@@ -191,7 +213,7 @@ class AbstractOperations(util.ModuleClsProxy):
 
         """
 
-        def decorate(fn):
+        def decorate(fn: _C) -> _C:
             cls._to_impl.dispatch_for(op_cls)(fn)
             return fn
 
@@ -213,7 +235,7 @@ class AbstractOperations(util.ModuleClsProxy):
         table_name: str,
         schema: Optional[str] = None,
         recreate: Literal["auto", "always", "never"] = "auto",
-        partial_reordering: Optional[tuple] = None,
+        partial_reordering: Optional[Tuple[Any, ...]] = None,
         copy_from: Optional[Table] = None,
         table_args: Tuple[Any, ...] = (),
         table_kwargs: Mapping[str, Any] = util.immutabledict(),
@@ -382,6 +404,32 @@ class AbstractOperations(util.ModuleClsProxy):
 
         return self.migration_context
 
+    @overload
+    def invoke(self, operation: CreateTableOp) -> Table: ...
+
+    @overload
+    def invoke(
+        self,
+        operation: Union[
+            AddConstraintOp,
+            DropConstraintOp,
+            CreateIndexOp,
+            DropIndexOp,
+            AddColumnOp,
+            AlterColumnOp,
+            AlterTableOp,
+            CreateTableCommentOp,
+            DropTableCommentOp,
+            DropColumnOp,
+            BulkInsertOp,
+            DropTableOp,
+            ExecuteSQLOp,
+        ],
+    ) -> None: ...
+
+    @overload
+    def invoke(self, operation: MigrateOperation) -> Any: ...
+
     def invoke(self, operation: MigrateOperation) -> Any:
         """Given a :class:`.MigrateOperation`, invoke it in terms of
         this :class:`.Operations` instance.
@@ -416,7 +464,7 @@ class AbstractOperations(util.ModuleClsProxy):
         names will be converted along conventions.  If the ``target_metadata``
         contains the naming convention
         ``{"ck": "ck_bool_%(table_name)s_%(constraint_name)s"}``, then the
-        output of the following:
+        output of the following::
 
             op.add_column("t", "x", Boolean(name="x"))
 
@@ -570,6 +618,7 @@ class Operations(AbstractOperations):
             column: Column[Any],
             *,
             schema: Optional[str] = None,
+            if_not_exists: Optional[bool] = None,
         ) -> None:
             """Issue an "add column" instruction using the current
             migration context.
@@ -646,6 +695,10 @@ class Operations(AbstractOperations):
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param if_not_exists: If True, adds IF NOT EXISTS operator
+             when creating the new column for compatible dialects
+
+             .. versionadded:: 1.16.0
 
             """  # noqa: E501
             ...
@@ -657,12 +710,16 @@ class Operations(AbstractOperations):
             *,
             nullable: Optional[bool] = None,
             comment: Union[str, Literal[False], None] = False,
-            server_default: Any = False,
+            server_default: Union[
+                str, bool, Identity, Computed, TextClause, None
+            ] = False,
             new_column_name: Optional[str] = None,
-            type_: Union[TypeEngine, Type[TypeEngine], None] = None,
-            existing_type: Union[TypeEngine, Type[TypeEngine], None] = 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
+                str, bool, Identity, Computed, TextClause, None
             ] = False,
             existing_nullable: Optional[bool] = None,
             existing_comment: Optional[str] = None,
@@ -756,7 +813,7 @@ class Operations(AbstractOperations):
         def bulk_insert(
             self,
             table: Union[Table, TableClause],
-            rows: List[dict],
+            rows: List[Dict[str, Any]],
             *,
             multiinsert: bool = True,
         ) -> None:
@@ -1023,7 +1080,7 @@ class Operations(AbstractOperations):
             self,
             index_name: Optional[str],
             table_name: str,
-            columns: Sequence[Union[str, TextClause, Function[Any]]],
+            columns: Sequence[Union[str, TextClause, ColumnElement[Any]]],
             *,
             schema: Optional[str] = None,
             unique: bool = False,
@@ -1124,7 +1181,11 @@ class Operations(AbstractOperations):
             ...
 
         def create_table(
-            self, table_name: str, *columns: SchemaItem, **kw: Any
+            self,
+            table_name: str,
+            *columns: SchemaItem,
+            if_not_exists: Optional[bool] = None,
+            **kw: Any,
         ) -> Table:
             r"""Issue a "create table" instruction using the current migration
             context.
@@ -1196,6 +1257,10 @@ class Operations(AbstractOperations):
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param if_not_exists: If True, adds IF NOT EXISTS operator when
+             creating the new table.
+
+             .. versionadded:: 1.13.3
             :param \**kw: Other keyword arguments are passed to the underlying
              :class:`sqlalchemy.schema.Table` object created for the command.
 
@@ -1301,6 +1366,11 @@ class Operations(AbstractOperations):
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param if_exists: If True, adds IF EXISTS operator when
+             dropping the new column for compatible dialects
+
+             .. versionadded:: 1.16.0
+
             :param mssql_drop_check: Optional boolean.  When ``True``, on
              Microsoft SQL Server only, first
              drop the CHECK constraint on the column using a
@@ -1322,7 +1392,6 @@ class Operations(AbstractOperations):
              then exec's a separate DROP CONSTRAINT for that default.  Only
              works if the column has exactly one FK constraint which refers to
              it, at the moment.
-
             """  # noqa: E501
             ...
 
@@ -1333,6 +1402,7 @@ class Operations(AbstractOperations):
             type_: Optional[str] = None,
             *,
             schema: Optional[str] = None,
+            if_exists: Optional[bool] = None,
         ) -> None:
             r"""Drop a constraint of the given name, typically via DROP CONSTRAINT.
 
@@ -1344,6 +1414,10 @@ class Operations(AbstractOperations):
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param if_exists: If True, adds IF EXISTS operator when
+             dropping the constraint
+
+             .. versionadded:: 1.16.0
 
             """  # noqa: E501
             ...
@@ -1387,7 +1461,12 @@ class Operations(AbstractOperations):
             ...
 
         def drop_table(
-            self, table_name: str, *, schema: Optional[str] = None, **kw: Any
+            self,
+            table_name: str,
+            *,
+            schema: Optional[str] = None,
+            if_exists: Optional[bool] = None,
+            **kw: Any,
         ) -> None:
             r"""Issue a "drop table" instruction using the current
             migration context.
@@ -1402,6 +1481,10 @@ class Operations(AbstractOperations):
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param if_exists: If True, adds IF EXISTS operator when
+             dropping the table.
+
+             .. versionadded:: 1.13.3
             :param \**kw: Other keyword arguments are passed to the underlying
              :class:`sqlalchemy.schema.Table` object created for the command.
 
@@ -1560,7 +1643,7 @@ class BatchOperations(AbstractOperations):
 
     impl: BatchOperationsImpl
 
-    def _noop(self, operation):
+    def _noop(self, operation: Any) -> NoReturn:
         raise NotImplementedError(
             "The %s method does not apply to a batch table alter operation."
             % operation
@@ -1577,6 +1660,7 @@ class BatchOperations(AbstractOperations):
             *,
             insert_before: Optional[str] = None,
             insert_after: Optional[str] = None,
+            if_not_exists: Optional[bool] = None,
         ) -> None:
             """Issue an "add column" instruction using the current
             batch migration context.
@@ -1596,8 +1680,10 @@ class BatchOperations(AbstractOperations):
             comment: Union[str, Literal[False], None] = False,
             server_default: Any = False,
             new_column_name: Optional[str] = None,
-            type_: Union[TypeEngine, Type[TypeEngine], None] = None,
-            existing_type: Union[TypeEngine, Type[TypeEngine], None] = 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
             ] = False,
@@ -1652,7 +1738,7 @@ class BatchOperations(AbstractOperations):
 
         def create_exclude_constraint(
             self, constraint_name: str, *elements: Any, **kw: Any
-        ):
+        ) -> Optional[Table]:
             """Issue a "create exclude constraint" instruction using the
             current batch migration context.
 
@@ -1668,7 +1754,7 @@ class BatchOperations(AbstractOperations):
 
         def create_foreign_key(
             self,
-            constraint_name: str,
+            constraint_name: Optional[str],
             referent_table: str,
             local_cols: List[str],
             remote_cols: List[str],
@@ -1718,7 +1804,7 @@ class BatchOperations(AbstractOperations):
             ...
 
         def create_primary_key(
-            self, constraint_name: str, columns: List[str]
+            self, constraint_name: Optional[str], columns: List[str]
         ) -> None:
             """Issue a "create primary key" instruction using the
             current batch migration context.