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

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,7 +904,7 @@ 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,
             index.expressions,
             schema=index.table.schema,
@@ -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

alembic/runtime/environment.py

@@ -228,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.
 
@@ -339,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]:
@@ -367,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::
@@ -401,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(
@@ -941,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

alembic/script/base.py

@@ -41,7 +41,7 @@ try:
         from zoneinfo import ZoneInfoNotFoundError
     else:
         from backports.zoneinfo import ZoneInfo  # type: ignore[import-not-found,no-redef] # noqa: E501
-        from backports.zoneinfo import ZoneInfoNotFoundError  # type: ignore[import-not-found,no-redef] # noqa: E501
+        from backports.zoneinfo import ZoneInfoNotFoundError  # type: ignore[no-redef] # noqa: E501
 except ImportError:
     ZoneInfo = None  # type: ignore[assignment, misc]
 
@@ -119,7 +119,7 @@ class ScriptDirectory:
             return loc[0]
 
     @util.memoized_property
-    def _version_locations(self):
+    def _version_locations(self) -> Sequence[str]:
         if self.version_locations:
             return [
                 os.path.abspath(util.coerce_resource_to_filename(location))
@@ -303,24 +303,22 @@ class ScriptDirectory:
             ):
                 yield cast(Script, rev)
 
-    def get_revisions(self, id_: _GetRevArg) -> Tuple[Optional[Script], ...]:
+    def get_revisions(self, id_: _GetRevArg) -> Tuple[Script, ...]:
         """Return the :class:`.Script` instance with the given rev identifier,
         symbolic name, or sequence of identifiers.
 
         """
         with self._catch_revision_errors():
             return cast(
-                Tuple[Optional[Script], ...],
+                Tuple[Script, ...],
                 self.revision_map.get_revisions(id_),
             )
 
-    def get_all_current(self, id_: Tuple[str, ...]) -> Set[Optional[Script]]:
+    def get_all_current(self, id_: Tuple[str, ...]) -> Set[Script]:
         with self._catch_revision_errors():
-            return cast(
-                Set[Optional[Script]], self.revision_map._get_all_current(id_)
-            )
+            return cast(Set[Script], self.revision_map._get_all_current(id_))
 
-    def get_revision(self, id_: str) -> Optional[Script]:
+    def get_revision(self, id_: str) -> Script:
         """Return the :class:`.Script` instance with the given rev id.
 
         .. seealso::
@@ -330,7 +328,7 @@ class ScriptDirectory:
         """
 
         with self._catch_revision_errors():
-            return cast(Optional[Script], self.revision_map.get_revision(id_))
+            return cast(Script, self.revision_map.get_revision(id_))
 
     def as_revision_number(
         self, id_: Optional[str]
@@ -585,7 +583,7 @@ class ScriptDirectory:
         util.load_python_file(self.dir, "env.py")
 
     @property
-    def env_py_location(self):
+    def env_py_location(self) -> str:
         return os.path.abspath(os.path.join(self.dir, "env.py"))
 
     def _generate_template(self, src: str, dest: str, **kw: Any) -> None:
@@ -684,7 +682,7 @@ class ScriptDirectory:
                 self.revision_map.get_revisions(head),
             )
             for h in heads:
-                assert h != "base"
+                assert h != "base"  # type: ignore[comparison-overlap]
 
         if len(set(heads)) != len(heads):
             raise util.CommandError("Duplicate head revisions specified")
@@ -823,7 +821,7 @@ class Script(revision.Revision):
         self.path = path
         super().__init__(
             rev_id,
-            module.down_revision,  # type: ignore[attr-defined]
+            module.down_revision,
             branch_labels=util.to_tuple(
                 getattr(module, "branch_labels", None), default=()
             ),
@@ -856,7 +854,7 @@ class Script(revision.Revision):
         if doc:
             if hasattr(self.module, "_alembic_source_encoding"):
                 doc = doc.decode(  # type: ignore[attr-defined]
-                    self.module._alembic_source_encoding  # type: ignore[attr-defined] # noqa
+                    self.module._alembic_source_encoding
                 )
             return doc.strip()  # type: ignore[union-attr]
         else:
@@ -898,7 +896,7 @@ class Script(revision.Revision):
         )
         return entry
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "%s -> %s%s%s%s, %s" % (
             self._format_down_revision(),
             self.revision,