API refactor

venv/lib/python3.12/site-packages/alembic/script/revision.py

@@ -14,6 +14,7 @@ from typing import Iterator
 from typing import List
 from typing import Optional
 from typing import overload
+from typing import Protocol
 from typing import Sequence
 from typing import Set
 from typing import Tuple
@@ -47,6 +48,17 @@ _relative_destination = re.compile(r"(?:(.+?)@)?(\w+)?((?:\+|-)\d+)")
 _revision_illegal_chars = ["@", "-", "+"]
 
 
+class _CollectRevisionsProtocol(Protocol):
+    def __call__(
+        self,
+        upper: _RevisionIdentifierType,
+        lower: _RevisionIdentifierType,
+        inclusive: bool,
+        implicit_base: bool,
+        assert_relative_length: bool,
+    ) -> Tuple[Set[Revision], Tuple[Optional[_RevisionOrBase], ...]]: ...
+
+
 class RevisionError(Exception):
     pass
 
@@ -396,7 +408,7 @@ class RevisionMap:
                 for rev in self._get_ancestor_nodes(
                     [revision],
                     include_dependencies=False,
-                    map_=cast(_RevisionMapType, map_),
+                    map_=map_,
                 ):
                     if rev is revision:
                         continue
@@ -707,9 +719,11 @@ class RevisionMap:
             resolved_target = target
 
         resolved_test_against_revs = [
-            self._revision_for_ident(test_against_rev)
-            if not isinstance(test_against_rev, Revision)
-            else test_against_rev
+            (
+                self._revision_for_ident(test_against_rev)
+                if not isinstance(test_against_rev, Revision)
+                else test_against_rev
+            )
             for test_against_rev in util.to_tuple(
                 test_against_revs, default=()
             )
@@ -791,7 +805,7 @@ class RevisionMap:
         The iterator yields :class:`.Revision` objects.
 
         """
-        fn: Callable
+        fn: _CollectRevisionsProtocol
         if select_for_downgrade:
             fn = self._collect_downgrade_revisions
         else:
@@ -818,7 +832,7 @@ class RevisionMap:
     ) -> Iterator[Any]:
         if omit_immediate_dependencies:
 
-            def fn(rev):
+            def fn(rev: Revision) -> Iterable[str]:
                 if rev not in targets:
                     return rev._all_nextrev
                 else:
@@ -826,12 +840,12 @@ class RevisionMap:
 
         elif include_dependencies:
 
-            def fn(rev):
+            def fn(rev: Revision) -> Iterable[str]:
                 return rev._all_nextrev
 
         else:
 
-            def fn(rev):
+            def fn(rev: Revision) -> Iterable[str]:
                 return rev.nextrev
 
         return self._iterate_related_revisions(
@@ -847,12 +861,12 @@ class RevisionMap:
     ) -> Iterator[Revision]:
         if include_dependencies:
 
-            def fn(rev):
+            def fn(rev: Revision) -> Iterable[str]:
                 return rev._normalized_down_revisions
 
         else:
 
-            def fn(rev):
+            def fn(rev: Revision) -> Iterable[str]:
                 return rev._versioned_down_revisions
 
         return self._iterate_related_revisions(
@@ -861,7 +875,7 @@ class RevisionMap:
 
     def _iterate_related_revisions(
         self,
-        fn: Callable,
+        fn: Callable[[Revision], Iterable[str]],
         targets: Collection[Optional[_RevisionOrBase]],
         map_: Optional[_RevisionMapType],
         check: bool = False,
@@ -923,7 +937,7 @@ class RevisionMap:
 
         id_to_rev = self._revision_map
 
-        def get_ancestors(rev_id):
+        def get_ancestors(rev_id: str) -> Set[str]:
             return {
                 r.revision
                 for r in self._get_ancestor_nodes([id_to_rev[rev_id]])
@@ -1003,9 +1017,9 @@ class RevisionMap:
                         # each time but it was getting complicated
                         current_heads[current_candidate_idx] = heads_to_add[0]
                         current_heads.extend(heads_to_add[1:])
-                        ancestors_by_idx[
-                            current_candidate_idx
-                        ] = get_ancestors(heads_to_add[0])
+                        ancestors_by_idx[current_candidate_idx] = (
+                            get_ancestors(heads_to_add[0])
+                        )
                         ancestors_by_idx.extend(
                             get_ancestors(head) for head in heads_to_add[1:]
                         )
@@ -1041,7 +1055,7 @@ class RevisionMap:
         children: Sequence[Optional[_RevisionOrBase]]
         for _ in range(abs(steps)):
             if steps > 0:
-                assert initial != "base"
+                assert initial != "base"  # type: ignore[comparison-overlap]
                 # Walk up
                 walk_up = [
                     is_revision(rev)
@@ -1055,7 +1069,7 @@ class RevisionMap:
                     children = walk_up
             else:
                 # Walk down
-                if initial == "base":
+                if initial == "base":  # type: ignore[comparison-overlap]
                     children = ()
                 else:
                     children = self.get_revisions(
@@ -1170,9 +1184,13 @@ class RevisionMap:
                         branch_label = symbol
                 # Walk down the tree to find downgrade target.
                 rev = self._walk(
-                    start=self.get_revision(symbol)
-                    if branch_label is None
-                    else self.get_revision("%s@%s" % (branch_label, symbol)),
+                    start=(
+                        self.get_revision(symbol)
+                        if branch_label is None
+                        else self.get_revision(
+                            "%s@%s" % (branch_label, symbol)
+                        )
+                    ),
                     steps=rel_int,
                     no_overwalk=assert_relative_length,
                 )
@@ -1189,7 +1207,7 @@ class RevisionMap:
         # No relative destination given, revision specified is absolute.
         branch_label, _, symbol = target.rpartition("@")
         if not branch_label:
-            branch_label = None  # type:ignore[assignment]
+            branch_label = None
         return branch_label, self.get_revision(symbol)
 
     def _parse_upgrade_target(
@@ -1290,9 +1308,13 @@ class RevisionMap:
                 )
             return (
                 self._walk(
-                    start=self.get_revision(symbol)
-                    if branch_label is None
-                    else self.get_revision("%s@%s" % (branch_label, symbol)),
+                    start=(
+                        self.get_revision(symbol)
+                        if branch_label is None
+                        else self.get_revision(
+                            "%s@%s" % (branch_label, symbol)
+                        )
+                    ),
                     steps=relative,
                     no_overwalk=assert_relative_length,
                 ),
@@ -1301,11 +1323,11 @@ class RevisionMap:
     def _collect_downgrade_revisions(
         self,
         upper: _RevisionIdentifierType,
-        target: _RevisionIdentifierType,
+        lower: _RevisionIdentifierType,
         inclusive: bool,
         implicit_base: bool,
         assert_relative_length: bool,
-    ) -> Any:
+    ) -> Tuple[Set[Revision], Tuple[Optional[_RevisionOrBase], ...]]:
         """
         Compute the set of current revisions specified by :upper, and the
         downgrade target specified by :target. Return all dependents of target
@@ -1316,7 +1338,7 @@ class RevisionMap:
 
         branch_label, target_revision = self._parse_downgrade_target(
             current_revisions=upper,
-            target=target,
+            target=lower,
             assert_relative_length=assert_relative_length,
         )
         if target_revision == "base":
@@ -1408,7 +1430,7 @@ class RevisionMap:
         inclusive: bool,
         implicit_base: bool,
         assert_relative_length: bool,
-    ) -> Tuple[Set[Revision], Tuple[Optional[_RevisionOrBase]]]:
+    ) -> Tuple[Set[Revision], Tuple[Revision, ...]]:
         """
         Compute the set of required revisions specified by :upper, and the
         current set of active revisions specified by :lower. Find the
@@ -1500,7 +1522,7 @@ class RevisionMap:
             )
             needs.intersection_update(lower_descendents)
 
-        return needs, tuple(targets)  # type:ignore[return-value]
+        return needs, tuple(targets)
 
     def _get_all_current(
         self, id_: Tuple[str, ...]
@@ -1681,15 +1703,13 @@ class Revision:
 
 
 @overload
-def tuple_rev_as_scalar(rev: None) -> None:
-    ...
+def tuple_rev_as_scalar(rev: None) -> None: ...
 
 
 @overload
 def tuple_rev_as_scalar(
-    rev: Union[Tuple[_T, ...], List[_T]]
-) -> Union[_T, Tuple[_T, ...], List[_T]]:
-    ...
+    rev: Union[Tuple[_T, ...], List[_T]],
+) -> Union[_T, Tuple[_T, ...], List[_T]]: ...
 
 
 def tuple_rev_as_scalar(