API refactor

venv/lib/python3.12/site-packages/alembic/op.pyi

@@ -12,6 +12,7 @@ from typing import List
 from typing import Literal
 from typing import Mapping
 from typing import Optional
+from typing import overload
 from typing import Sequence
 from typing import Tuple
 from typing import Type
@@ -26,7 +27,6 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import conv
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.expression import TableClause
-    from sqlalchemy.sql.functions import Function
     from sqlalchemy.sql.schema import Column
     from sqlalchemy.sql.schema import Computed
     from sqlalchemy.sql.schema import Identity
@@ -35,16 +35,36 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.type_api import TypeEngine
     from sqlalchemy.util import immutabledict
 
-    from .operations.ops import BatchOperations
+    from .operations.base import BatchOperations
+    from .operations.ops import AddColumnOp
+    from .operations.ops import AddConstraintOp
+    from .operations.ops import AlterColumnOp
+    from .operations.ops import AlterTableOp
+    from .operations.ops import BulkInsertOp
+    from .operations.ops import CreateIndexOp
+    from .operations.ops import CreateTableCommentOp
+    from .operations.ops import CreateTableOp
+    from .operations.ops import DropColumnOp
+    from .operations.ops import DropConstraintOp
+    from .operations.ops import DropIndexOp
+    from .operations.ops import DropTableCommentOp
+    from .operations.ops import DropTableOp
+    from .operations.ops import ExecuteSQLOp
     from .operations.ops import MigrateOperation
     from .runtime.migration import MigrationContext
     from .util.sqla_compat import _literal_bindparam
 
 _T = TypeVar("_T")
+_C = TypeVar("_C", bound=Callable[..., Any])
+
 ### end imports ###
 
 def add_column(
-    table_name: str, column: Column[Any], *, schema: Optional[str] = None
+    table_name: str,
+    column: Column[Any],
+    *,
+    schema: Optional[str] = None,
+    if_not_exists: Optional[bool] = None,
 ) -> None:
     """Issue an "add column" instruction using the current
     migration context.
@@ -121,6 +141,10 @@ def add_column(
      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
 
     """
 
@@ -130,12 +154,14 @@ def alter_column(
     *,
     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,
@@ -230,7 +256,7 @@ def batch_alter_table(
     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] = immutabledict({}),
@@ -377,7 +403,7 @@ def batch_alter_table(
 
 def bulk_insert(
     table: Union[Table, TableClause],
-    rows: List[dict],
+    rows: List[Dict[str, Any]],
     *,
     multiinsert: bool = True,
 ) -> None:
@@ -633,7 +659,7 @@ def create_foreign_key(
 def create_index(
     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,
@@ -730,7 +756,12 @@ def create_primary_key(
 
     """
 
-def create_table(table_name: str, *columns: SchemaItem, **kw: Any) -> Table:
+def create_table(
+    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.
 
@@ -801,6 +832,10 @@ def create_table(table_name: str, *columns: SchemaItem, **kw: Any) -> Table:
      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.
 
@@ -900,6 +935,11 @@ def drop_column(
      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
@@ -921,7 +961,6 @@ def drop_column(
      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.
-
     """
 
 def drop_constraint(
@@ -930,6 +969,7 @@ def drop_constraint(
     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.
 
@@ -941,6 +981,10 @@ def drop_constraint(
      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
 
     """
 
@@ -981,7 +1025,11 @@ def drop_index(
     """
 
 def drop_table(
-    table_name: str, *, schema: Optional[str] = None, **kw: Any
+    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.
@@ -996,6 +1044,10 @@ def drop_table(
      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.
 
@@ -1132,7 +1184,7 @@ def f(name: str) -> conv:
     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"))
 
@@ -1162,7 +1214,7 @@ def get_context() -> MigrationContext:
 
     """
 
-def implementation_for(op_cls: Any) -> Callable[..., Any]:
+def implementation_for(op_cls: Any) -> Callable[[_C], _C]:
     """Register an implementation for a given :class:`.MigrateOperation`.
 
     This is part of the operation extensibility API.
@@ -1174,7 +1226,7 @@ def implementation_for(op_cls: Any) -> Callable[..., Any]:
     """
 
 def inline_literal(
-    value: Union[str, int], type_: Optional[TypeEngine] = None
+    value: Union[str, int], type_: Optional[TypeEngine[Any]] = None
 ) -> _literal_bindparam:
     r"""Produce an 'inline literal' expression, suitable for
     using in an INSERT, UPDATE, or DELETE statement.
@@ -1218,6 +1270,27 @@ def inline_literal(
 
     """
 
+@overload
+def invoke(operation: CreateTableOp) -> Table: ...
+@overload
+def invoke(
+    operation: Union[
+        AddConstraintOp,
+        DropConstraintOp,
+        CreateIndexOp,
+        DropIndexOp,
+        AddColumnOp,
+        AlterColumnOp,
+        AlterTableOp,
+        CreateTableCommentOp,
+        DropTableCommentOp,
+        DropColumnOp,
+        BulkInsertOp,
+        DropTableOp,
+        ExecuteSQLOp,
+    ],
+) -> None: ...
+@overload
 def invoke(operation: MigrateOperation) -> Any:
     """Given a :class:`.MigrateOperation`, invoke it in terms of
     this :class:`.Operations` instance.
@@ -1226,7 +1299,7 @@ def invoke(operation: MigrateOperation) -> Any:
 
 def register_operation(
     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