main commit

venv/lib/python3.12/site-packages/alembic/autogenerate/rewriter.py

@@ -4,7 +4,7 @@ from typing import Any
 from typing import Callable
 from typing import Iterator
 from typing import List
-from typing import Tuple
+from typing import Optional
 from typing import Type
 from typing import TYPE_CHECKING
 from typing import Union
@@ -16,18 +16,12 @@ if TYPE_CHECKING:
     from ..operations.ops import AddColumnOp
     from ..operations.ops import AlterColumnOp
     from ..operations.ops import CreateTableOp
-    from ..operations.ops import DowngradeOps
     from ..operations.ops import MigrateOperation
     from ..operations.ops import MigrationScript
     from ..operations.ops import ModifyTableOps
     from ..operations.ops import OpContainer
-    from ..operations.ops import UpgradeOps
+    from ..runtime.environment import _GetRevArg
     from ..runtime.migration import MigrationContext
-    from ..script.revision import _GetRevArg
-
-ProcessRevisionDirectiveFn = Callable[
-    ["MigrationContext", "_GetRevArg", List["MigrationScript"]], None
-]
 
 
 class Rewriter:
@@ -58,21 +52,15 @@ class Rewriter:
 
     _traverse = util.Dispatcher()
 
-    _chained: Tuple[Union[ProcessRevisionDirectiveFn, Rewriter], ...] = ()
+    _chained: Optional[Rewriter] = None
 
     def __init__(self) -> None:
         self.dispatch = util.Dispatcher()
 
-    def chain(
-        self,
-        other: Union[
-            ProcessRevisionDirectiveFn,
-            Rewriter,
-        ],
-    ) -> Rewriter:
+    def chain(self, other: Rewriter) -> Rewriter:
         """Produce a "chain" of this :class:`.Rewriter` to another.
 
-        This allows two or more rewriters to operate serially on a stream,
+        This allows two rewriters to operate serially on a stream,
         e.g.::
 
             writer1 = autogenerate.Rewriter()
@@ -101,7 +89,7 @@ class Rewriter:
         """
         wr = self.__class__.__new__(self.__class__)
         wr.__dict__.update(self.__dict__)
-        wr._chained += (other,)
+        wr._chained = other
         return wr
 
     def rewrites(
@@ -113,7 +101,7 @@ class Rewriter:
             Type[CreateTableOp],
             Type[ModifyTableOps],
         ],
-    ) -> Callable[..., Any]:
+    ) -> Callable:
         """Register a function as rewriter for a given type.
 
         The function should receive three arguments, which are
@@ -158,8 +146,8 @@ class Rewriter:
         directives: List[MigrationScript],
     ) -> None:
         self.process_revision_directives(context, revision, directives)
-        for process_revision_directives in self._chained:
-            process_revision_directives(context, revision, directives)
+        if self._chained:
+            self._chained(context, revision, directives)
 
     @_traverse.dispatch_for(ops.MigrationScript)
     def _traverse_script(
@@ -168,7 +156,7 @@ class Rewriter:
         revision: _GetRevArg,
         directive: MigrationScript,
     ) -> None:
-        upgrade_ops_list: List[UpgradeOps] = []
+        upgrade_ops_list = []
         for upgrade_ops in directive.upgrade_ops_list:
             ret = self._traverse_for(context, revision, upgrade_ops)
             if len(ret) != 1:
@@ -176,10 +164,9 @@ class Rewriter:
                     "Can only return single object for UpgradeOps traverse"
                 )
             upgrade_ops_list.append(ret[0])
-
         directive.upgrade_ops = upgrade_ops_list
 
-        downgrade_ops_list: List[DowngradeOps] = []
+        downgrade_ops_list = []
         for downgrade_ops in directive.downgrade_ops_list:
             ret = self._traverse_for(context, revision, downgrade_ops)
             if len(ret) != 1: