alembic 1.12.1__py3-none-any.whl → 1.13.1__py3-none-any.whl

alembic/operations/batch.py

@@ -1,3 +1,6 @@
+# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
+# mypy: no-warn-return-any, allow-any-generics
+
 from __future__ import annotations
 
 from typing import Any
@@ -17,7 +20,7 @@ from sqlalchemy import PrimaryKeyConstraint
 from sqlalchemy import schema as sql_schema
 from sqlalchemy import Table
 from sqlalchemy import types as sqltypes
-from sqlalchemy.events import SchemaEventTarget
+from sqlalchemy.sql.schema import SchemaEventTarget
 from sqlalchemy.util import OrderedDict
 from sqlalchemy.util import topological
 
@@ -374,7 +377,7 @@ class ApplyBatchImpl:
         for idx_existing in self.indexes.values():
             # this is a lift-and-move from Table.to_metadata
 
-            if idx_existing._column_flag:  # type: ignore
+            if idx_existing._column_flag:
                 continue
 
             idx_copy = Index(
@@ -403,9 +406,7 @@ class ApplyBatchImpl:
     def _setup_referent(
         self, metadata: MetaData, constraint: ForeignKeyConstraint
     ) -> None:
-        spec = constraint.elements[
-            0
-        ]._get_colspec()  # type:ignore[attr-defined]
+        spec = constraint.elements[0]._get_colspec()
         parts = spec.split(".")
         tname = parts[-2]
         if len(parts) == 3:
@@ -546,9 +547,7 @@ class ApplyBatchImpl:
             else:
                 sql_schema.DefaultClause(
                     server_default  # type: ignore[arg-type]
-                )._set_parent(  # type:ignore[attr-defined]
-                    existing
-                )
+                )._set_parent(existing)
         if autoincrement is not None:
             existing.autoincrement = bool(autoincrement)
 

alembic/operations/ops.py

@@ -5,6 +5,7 @@ import re
 from typing import Any
 from typing import Callable
 from typing import cast
+from typing import Dict
 from typing import FrozenSet
 from typing import Iterator
 from typing import List
@@ -15,6 +16,7 @@ from typing import Set
 from typing import Tuple
 from typing import Type
 from typing import TYPE_CHECKING
+from typing import TypeVar
 from typing import Union
 
 from sqlalchemy.types import NULLTYPE
@@ -53,6 +55,9 @@ if TYPE_CHECKING:
     from ..runtime.migration import MigrationContext
     from ..script.revision import _RevIdType
 
+_T = TypeVar("_T", bound=Any)
+_AC = TypeVar("_AC", bound="AddConstraintOp")
+
 
 class MigrateOperation:
     """base class for migration command and organization objects.
@@ -70,7 +75,7 @@ class MigrateOperation:
     """
 
     @util.memoized_property
-    def info(self):
+    def info(self) -> Dict[Any, Any]:
         """A dictionary that may be used to store arbitrary information
         along with this :class:`.MigrateOperation` object.
 
@@ -92,12 +97,14 @@ class AddConstraintOp(MigrateOperation):
     add_constraint_ops = util.Dispatcher()
 
     @property
-    def constraint_type(self):
+    def constraint_type(self) -> str:
         raise NotImplementedError()
 
     @classmethod
-    def register_add_constraint(cls, type_: str) -> Callable:
-        def go(klass):
+    def register_add_constraint(
+        cls, type_: str
+    ) -> Callable[[Type[_AC]], Type[_AC]]:
+        def go(klass: Type[_AC]) -> Type[_AC]:
             cls.add_constraint_ops.dispatch_for(type_)(klass.from_constraint)
             return klass
 
@@ -105,7 +112,7 @@ class AddConstraintOp(MigrateOperation):
 
     @classmethod
     def from_constraint(cls, constraint: Constraint) -> AddConstraintOp:
-        return cls.add_constraint_ops.dispatch(constraint.__visit_name__)(
+        return cls.add_constraint_ops.dispatch(constraint.__visit_name__)(  # type: ignore[no-any-return]  # noqa: E501
             constraint
         )
 
@@ -398,7 +405,7 @@ class CreateUniqueConstraintOp(AddConstraintOp):
 
         uq_constraint = cast("UniqueConstraint", constraint)
 
-        kw: dict = {}
+        kw: Dict[str, Any] = {}
         if uq_constraint.deferrable:
             kw["deferrable"] = uq_constraint.deferrable
         if uq_constraint.initially:
@@ -532,7 +539,7 @@ class CreateForeignKeyOp(AddConstraintOp):
     @classmethod
     def from_constraint(cls, constraint: Constraint) -> CreateForeignKeyOp:
         fk_constraint = cast("ForeignKeyConstraint", constraint)
-        kw: dict = {}
+        kw: Dict[str, Any] = {}
         if fk_constraint.onupdate:
             kw["onupdate"] = fk_constraint.onupdate
         if fk_constraint.ondelete:
@@ -897,9 +904,9 @@ class CreateIndexOp(MigrateOperation):
     def from_index(cls, index: Index) -> CreateIndexOp:
         assert index.table is not None
         return cls(
-            index.name,  # type: ignore[arg-type]
+            index.name,
             index.table.name,
-            sqla_compat._get_index_expressions(index),
+            index.expressions,
             schema=index.table.schema,
             unique=index.unique,
             **index.kwargs,
@@ -1054,6 +1061,7 @@ class DropIndexOp(MigrateOperation):
             table_name=index.table.name,
             schema=index.table.schema,
             _reverse=CreateIndexOp.from_index(index),
+            unique=index.unique,
             **index.kwargs,
         )
 
@@ -1182,7 +1190,7 @@ class CreateTableOp(MigrateOperation):
 
         return cls(
             table.name,
-            list(table.c) + list(table.constraints),  # type:ignore[arg-type]
+            list(table.c) + list(table.constraints),
             schema=table.schema,
             _namespace_metadata=_namespace_metadata,
             # given a Table() object, this Table will contain full Index()
@@ -1534,7 +1542,7 @@ class CreateTableCommentOp(AlterTableOp):
         )
         return operations.invoke(op)
 
-    def reverse(self):
+    def reverse(self) -> Union[CreateTableCommentOp, DropTableCommentOp]:
         """Reverses the COMMENT ON operation against a table."""
         if self.existing_comment is None:
             return DropTableCommentOp(
@@ -1550,14 +1558,16 @@ class CreateTableCommentOp(AlterTableOp):
                 schema=self.schema,
             )
 
-    def to_table(self, migration_context=None):
+    def to_table(
+        self, migration_context: Optional[MigrationContext] = None
+    ) -> Table:
         schema_obj = schemaobj.SchemaObjects(migration_context)
 
         return schema_obj.table(
             self.table_name, schema=self.schema, comment=self.comment
         )
 
-    def to_diff_tuple(self):
+    def to_diff_tuple(self) -> Tuple[Any, ...]:
         return ("add_table_comment", self.to_table(), self.existing_comment)
 
 
@@ -1629,18 +1639,20 @@ class DropTableCommentOp(AlterTableOp):
         )
         return operations.invoke(op)
 
-    def reverse(self):
+    def reverse(self) -> CreateTableCommentOp:
         """Reverses the COMMENT ON operation against a table."""
         return CreateTableCommentOp(
             self.table_name, self.existing_comment, schema=self.schema
         )
 
-    def to_table(self, migration_context=None):
+    def to_table(
+        self, migration_context: Optional[MigrationContext] = None
+    ) -> Table:
         schema_obj = schemaobj.SchemaObjects(migration_context)
 
         return schema_obj.table(self.table_name, schema=self.schema)
 
-    def to_diff_tuple(self):
+    def to_diff_tuple(self) -> Tuple[Any, ...]:
         return ("remove_table_comment", self.to_table())
 
 
@@ -1817,8 +1829,10 @@ class AlterColumnOp(AlterTableOp):
         comment: Optional[Union[str, Literal[False]]] = False,
         server_default: Any = False,
         new_column_name: Optional[str] = None,
-        type_: Optional[Union[TypeEngine, Type[TypeEngine]]] = None,
-        existing_type: Optional[Union[TypeEngine, Type[TypeEngine]]] = 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]
         ] = False,
@@ -1938,8 +1952,10 @@ class AlterColumnOp(AlterTableOp):
         comment: Optional[Union[str, Literal[False]]] = False,
         server_default: Any = False,
         new_column_name: Optional[str] = None,
-        type_: Optional[Union[TypeEngine, Type[TypeEngine]]] = None,
-        existing_type: Optional[Union[TypeEngine, Type[TypeEngine]]] = 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]
         ] = False,
@@ -2019,11 +2035,11 @@ class AddColumnOp(AlterTableOp):
     ) -> Tuple[str, Optional[str], str, Column[Any]]:
         return ("add_column", self.schema, self.table_name, self.column)
 
-    def to_column(self) -> Column:
+    def to_column(self) -> Column[Any]:
         return self.column
 
     @classmethod
-    def from_column(cls, col: Column) -> AddColumnOp:
+    def from_column(cls, col: Column[Any]) -> AddColumnOp:
         return cls(col.table.name, col, schema=col.table.schema)
 
     @classmethod
@@ -2214,7 +2230,7 @@ class DropColumnOp(AlterTableOp):
 
     def to_column(
         self, migration_context: Optional[MigrationContext] = None
-    ) -> Column:
+    ) -> Column[Any]:
         if self._reverse is not None:
             return self._reverse.column
         schema_obj = schemaobj.SchemaObjects(migration_context)
@@ -2298,7 +2314,7 @@ class BulkInsertOp(MigrateOperation):
     def __init__(
         self,
         table: Union[Table, TableClause],
-        rows: List[dict],
+        rows: List[Dict[str, Any]],
         *,
         multiinsert: bool = True,
     ) -> None:
@@ -2311,7 +2327,7 @@ class BulkInsertOp(MigrateOperation):
         cls,
         operations: Operations,
         table: Union[Table, TableClause],
-        rows: List[dict],
+        rows: List[Dict[str, Any]],
         *,
         multiinsert: bool = True,
     ) -> None:
@@ -2607,7 +2623,7 @@ class UpgradeOps(OpContainer):
         self.upgrade_token = upgrade_token
 
     def reverse_into(self, downgrade_ops: DowngradeOps) -> DowngradeOps:
-        downgrade_ops.ops[:] = list(  # type:ignore[index]
+        downgrade_ops.ops[:] = list(
             reversed([op.reverse() for op in self.ops])
         )
         return downgrade_ops
@@ -2634,7 +2650,7 @@ class DowngradeOps(OpContainer):
         super().__init__(ops=ops)
         self.downgrade_token = downgrade_token
 
-    def reverse(self):
+    def reverse(self) -> UpgradeOps:
         return UpgradeOps(
             ops=list(reversed([op.reverse() for op in self.ops]))
         )
@@ -2665,6 +2681,8 @@ class MigrationScript(MigrateOperation):
     """
 
     _needs_render: Optional[bool]
+    _upgrade_ops: List[UpgradeOps]
+    _downgrade_ops: List[DowngradeOps]
 
     def __init__(
         self,
@@ -2692,7 +2710,7 @@ class MigrationScript(MigrateOperation):
         self.downgrade_ops = downgrade_ops
 
     @property
-    def upgrade_ops(self):
+    def upgrade_ops(self) -> Optional[UpgradeOps]:
         """An instance of :class:`.UpgradeOps`.
 
         .. seealso::
@@ -2711,13 +2729,15 @@ class MigrationScript(MigrateOperation):
             return self._upgrade_ops[0]
 
     @upgrade_ops.setter
-    def upgrade_ops(self, upgrade_ops):
+    def upgrade_ops(
+        self, upgrade_ops: Union[UpgradeOps, List[UpgradeOps]]
+    ) -> None:
         self._upgrade_ops = util.to_list(upgrade_ops)
         for elem in self._upgrade_ops:
             assert isinstance(elem, UpgradeOps)
 
     @property
-    def downgrade_ops(self):
+    def downgrade_ops(self) -> Optional[DowngradeOps]:
         """An instance of :class:`.DowngradeOps`.
 
         .. seealso::
@@ -2736,7 +2756,9 @@ class MigrationScript(MigrateOperation):
             return self._downgrade_ops[0]
 
     @downgrade_ops.setter
-    def downgrade_ops(self, downgrade_ops):
+    def downgrade_ops(
+        self, downgrade_ops: Union[DowngradeOps, List[DowngradeOps]]
+    ) -> None:
         self._downgrade_ops = util.to_list(downgrade_ops)
         for elem in self._downgrade_ops:
             assert isinstance(elem, DowngradeOps)

alembic/operations/schemaobj.py

@@ -1,3 +1,6 @@
+# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
+# mypy: no-warn-return-any, allow-any-generics
+
 from __future__ import annotations
 
 from typing import Any
@@ -274,10 +277,8 @@ class SchemaObjects:
         ForeignKey.
 
         """
-        if isinstance(fk._colspec, str):  # type:ignore[attr-defined]
-            table_key, cname = fk._colspec.rsplit(  # type:ignore[attr-defined]
-                ".", 1
-            )
+        if isinstance(fk._colspec, str):
+            table_key, cname = fk._colspec.rsplit(".", 1)
             sname, tname = self._parse_table_key(table_key)
             if table_key not in metadata.tables:
                 rel_t = sa_schema.Table(tname, metadata, schema=sname)

alembic/operations/toimpl.py

@@ -1,3 +1,6 @@
+# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
+# mypy: no-warn-return-any, allow-any-generics
+
 from typing import TYPE_CHECKING
 
 from sqlalchemy import schema as sa_schema
@@ -5,7 +8,7 @@ from sqlalchemy import schema as sa_schema
 from . import ops
 from .base import Operations
 from ..util.sqla_compat import _copy
-from ..util.sqla_compat import sqla_2
+from ..util.sqla_compat import sqla_14
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.schema import Table
@@ -98,8 +101,8 @@ def create_index(
     idx = operation.to_index(operations.migration_context)
     kw = {}
     if operation.if_not_exists is not None:
-        if not sqla_2:
-            raise NotImplementedError("SQLAlchemy 2.0+ required")
+        if not sqla_14:
+            raise NotImplementedError("SQLAlchemy 1.4+ required")
 
         kw["if_not_exists"] = operation.if_not_exists
     operations.impl.create_index(idx, **kw)
@@ -109,8 +112,8 @@ def create_index(
 def drop_index(operations: "Operations", operation: "ops.DropIndexOp") -> None:
     kw = {}
     if operation.if_exists is not None:
-        if not sqla_2:
-            raise NotImplementedError("SQLAlchemy 2.0+ required")
+        if not sqla_14:
+            raise NotImplementedError("SQLAlchemy 1.4+ required")
 
         kw["if_exists"] = operation.if_exists
 

alembic/runtime/environment.py

@@ -10,6 +10,7 @@ from typing import Mapping
 from typing import MutableMapping
 from typing import Optional
 from typing import overload
+from typing import Sequence
 from typing import TextIO
 from typing import Tuple
 from typing import TYPE_CHECKING
@@ -227,9 +228,9 @@ class EnvironmentContext(util.ModuleClsProxy):
         has been configured.
 
         """
-        return self.context_opts.get("as_sql", False)
+        return self.context_opts.get("as_sql", False)  # type: ignore[no-any-return]  # noqa: E501
 
-    def is_transactional_ddl(self):
+    def is_transactional_ddl(self) -> bool:
         """Return True if the context is configured to expect a
         transactional DDL capable backend.
 
@@ -338,7 +339,7 @@ class EnvironmentContext(util.ModuleClsProxy):
             line.
 
         """
-        return self.context_opts.get("tag", None)
+        return self.context_opts.get("tag", None)  # type: ignore[no-any-return]  # noqa: E501
 
     @overload
     def get_x_argument(self, as_dictionary: Literal[False]) -> List[str]:
@@ -366,7 +367,11 @@ class EnvironmentContext(util.ModuleClsProxy):
         The return value is a list, returned directly from the ``argparse``
         structure.  If ``as_dictionary=True`` is passed, the ``x`` arguments
         are parsed using ``key=value`` format into a dictionary that is
-        then returned.
+        then returned. If there is no ``=`` in the argument, value is an empty
+        string.
+
+        .. versionchanged:: 1.13.1 Support ``as_dictionary=True`` when
+           arguments are passed without the ``=`` symbol.
 
         For example, to support passing a database URL on the command line,
         the standard ``env.py`` script can be modified like this::
@@ -400,7 +405,12 @@ class EnvironmentContext(util.ModuleClsProxy):
         else:
             value = []
         if as_dictionary:
-            value = dict(arg.split("=", 1) for arg in value)
+            dict_value = {}
+            for arg in value:
+                x_key, _, x_value = arg.partition("=")
+                dict_value[x_key] = x_value
+            value = dict_value
+
         return value
 
     def configure(
@@ -416,7 +426,7 @@ class EnvironmentContext(util.ModuleClsProxy):
         tag: Optional[str] = None,
         template_args: Optional[Dict[str, Any]] = None,
         render_as_batch: bool = False,
-        target_metadata: Optional[MetaData] = None,
+        target_metadata: Union[MetaData, Sequence[MetaData], None] = None,
         include_name: Optional[IncludeNameFn] = None,
         include_object: Optional[IncludeObjectFn] = None,
         include_schemas: bool = False,
@@ -940,7 +950,7 @@ class EnvironmentContext(util.ModuleClsProxy):
     def execute(
         self,
         sql: Union[Executable, str],
-        execution_options: Optional[dict] = None,
+        execution_options: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Execute the given SQL using the current change context.
 

alembic/runtime/migration.py

@@ -1,3 +1,6 @@
+# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
+# mypy: no-warn-return-any, allow-any-generics
+
 from __future__ import annotations
 
 from contextlib import contextmanager
@@ -521,7 +524,7 @@ class MigrationContext:
                 start_from_rev = None
             elif start_from_rev is not None and self.script:
                 start_from_rev = [
-                    cast("Script", self.script.get_revision(sfr)).revision
+                    self.script.get_revision(sfr).revision
                     for sfr in util.to_list(start_from_rev)
                     if sfr not in (None, "base")
                 ]
@@ -652,7 +655,7 @@ class MigrationContext:
     def execute(
         self,
         sql: Union[Executable, str],
-        execution_options: Optional[dict] = None,
+        execution_options: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Execute a SQL construct or string statement.
 
@@ -1000,6 +1003,12 @@ class MigrationStep:
     is_upgrade: bool
     migration_fn: Any
 
+    if TYPE_CHECKING:
+
+        @property
+        def doc(self) -> Optional[str]:
+            ...
+
     @property
     def name(self) -> str:
         return self.migration_fn.__name__
@@ -1048,13 +1057,9 @@ class RevisionStep(MigrationStep):
         self.revision = revision
         self.is_upgrade = is_upgrade
         if is_upgrade:
-            self.migration_fn = (
-                revision.module.upgrade  # type:ignore[attr-defined]
-            )
+            self.migration_fn = revision.module.upgrade
         else:
-            self.migration_fn = (
-                revision.module.downgrade  # type:ignore[attr-defined]
-            )
+            self.migration_fn = revision.module.downgrade
 
     def __repr__(self):
         return "RevisionStep(%r, is_upgrade=%r)" % (
@@ -1070,7 +1075,7 @@ class RevisionStep(MigrationStep):
         )
 
     @property
-    def doc(self) -> str:
+    def doc(self) -> Optional[str]:
         return self.revision.doc
 
     @property
@@ -1168,7 +1173,18 @@ class RevisionStep(MigrationStep):
             }
             return tuple(set(self.to_revisions).difference(ancestors))
         else:
-            return self.to_revisions
+            # for each revision we plan to return, compute its ancestors
+            # (excluding self), and remove those from the final output since
+            # they are already accounted for.
+            ancestors = {
+                r.revision
+                for to_revision in self.to_revisions
+                for r in self.revision_map._get_ancestor_nodes(
+                    self.revision_map.get_revisions(to_revision), check=False
+                )
+                if r.revision != to_revision
+            }
+            return tuple(set(self.to_revisions).difference(ancestors))
 
     def unmerge_branch_idents(
         self, heads: Set[str]
@@ -1283,7 +1299,7 @@ class StampStep(MigrationStep):
     def __eq__(self, other):
         return (
             isinstance(other, StampStep)
-            and other.from_revisions == self.revisions
+            and other.from_revisions == self.from_revisions
             and other.to_revisions == self.to_revisions
             and other.branch_move == self.branch_move
             and self.is_upgrade == other.is_upgrade