API refactor

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