API refactor

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

@@ -17,6 +17,7 @@ from . import compare
 from . import render
 from .. import util
 from ..operations import ops
+from ..util import sqla_compat
 
 """Provide the 'autogenerate' feature which can produce migration operations
 automatically."""
@@ -27,6 +28,7 @@ if TYPE_CHECKING:
     from sqlalchemy.engine import Inspector
     from sqlalchemy.sql.schema import MetaData
     from sqlalchemy.sql.schema import SchemaItem
+    from sqlalchemy.sql.schema import Table
 
     from ..config import Config
     from ..operations.ops import DowngradeOps
@@ -164,6 +166,7 @@ def compare_metadata(context: MigrationContext, metadata: MetaData) -> Any:
     """
 
     migration_script = produce_migrations(context, metadata)
+    assert migration_script.upgrade_ops is not None
     return migration_script.upgrade_ops.as_diffs()
 
 
@@ -274,7 +277,7 @@ class AutogenContext:
     """Maintains configuration and state that's specific to an
     autogenerate operation."""
 
-    metadata: Optional[MetaData] = None
+    metadata: Union[MetaData, Sequence[MetaData], None] = None
     """The :class:`~sqlalchemy.schema.MetaData` object
     representing the destination.
 
@@ -329,8 +332,8 @@ class AutogenContext:
     def __init__(
         self,
         migration_context: MigrationContext,
-        metadata: Optional[MetaData] = None,
-        opts: Optional[dict] = None,
+        metadata: Union[MetaData, Sequence[MetaData], None] = None,
+        opts: Optional[Dict[str, Any]] = None,
         autogenerate: bool = True,
     ) -> None:
         if (
@@ -440,7 +443,7 @@ class AutogenContext:
     def run_object_filters(
         self,
         object_: SchemaItem,
-        name: Optional[str],
+        name: sqla_compat._ConstraintName,
         type_: NameFilterType,
         reflected: bool,
         compare_to: Optional[SchemaItem],
@@ -464,7 +467,7 @@ class AutogenContext:
     run_filters = run_object_filters
 
     @util.memoized_property
-    def sorted_tables(self):
+    def sorted_tables(self) -> List[Table]:
         """Return an aggregate of the :attr:`.MetaData.sorted_tables`
         collection(s).
 
@@ -480,7 +483,7 @@ class AutogenContext:
         return result
 
     @util.memoized_property
-    def table_key_to_table(self):
+    def table_key_to_table(self) -> Dict[str, Table]:
         """Return an aggregate  of the :attr:`.MetaData.tables` dictionaries.
 
         The :attr:`.MetaData.tables` collection is a dictionary of table key
@@ -491,7 +494,7 @@ class AutogenContext:
         objects contain the same table key, an exception is raised.
 
         """
-        result = {}
+        result: Dict[str, Table] = {}
         for m in util.to_list(self.metadata):
             intersect = set(result).intersection(set(m.tables))
             if intersect:
@@ -593,9 +596,9 @@ class RevisionContext:
         migration_script = self.generated_revisions[-1]
         if not getattr(migration_script, "_needs_render", False):
             migration_script.upgrade_ops_list[-1].upgrade_token = upgrade_token
-            migration_script.downgrade_ops_list[
-                -1
-            ].downgrade_token = downgrade_token
+            migration_script.downgrade_ops_list[-1].downgrade_token = (
+                downgrade_token
+            )
             migration_script._needs_render = True
         else:
             migration_script._upgrade_ops.append(