Various typing related updates

Change-Id: I778b63b1c438f31964d841576f0dd54ae1a5fadc

4 files changed, 29 insertions, 16 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index a5528ff..4ab8a35 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -1,5 +1,4 @@
-"""Provide the 'autogenerate' feature which can produce migration operations
-automatically."""
+from __future__ import annotations
 
 import contextlib
 from typing import Any
@@ -19,6 +18,9 @@ from . import render
 from .. import util
 from ..operations import ops
 
+"""Provide the 'autogenerate' feature which can produce migration operations
+automatically."""
+
 if TYPE_CHECKING:
     from sqlalchemy.engine import Connection
     from sqlalchemy.engine import Dialect
@@ -515,7 +517,7 @@ class RevisionContext:
             branch_labels=migration_script.branch_label,
             version_path=migration_script.version_path,
             depends_on=migration_script.depends_on,
-            **template_args
+            **template_args,
         )
 
     def run_autogenerate(
diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index 528b17a..693efae 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import logging
 import re
@@ -282,7 +284,7 @@ def _make_index(params: Dict[str, Any], conn_table: "Table") -> "Index":
         params["name"],
         *[conn_table.c[cname] for cname in params["column_names"]],
         unique=params["unique"],
-        _table=conn_table
+        _table=conn_table,
     )
     if "duplicates_constraint" in params:
         ix.info["duplicates_constraint"] = params["duplicates_constraint"]
@@ -294,7 +296,7 @@ def _make_unique_constraint(
 ) -> "UniqueConstraint":
     uq = sa_schema.UniqueConstraint(
         *[conn_table.c[cname] for cname in params["column_names"]],
-        name=params["name"]
+        name=params["name"],
     )
     if "duplicates_index" in params:
         uq.info["duplicates_index"] = params["duplicates_index"]
@@ -1245,7 +1247,7 @@ def _compare_foreign_keys(
         if isinstance(fk, sa_schema.ForeignKeyConstraint)
     )
 
-    conn_fks = [
+    conn_fks_list = [
         fk
         for fk in inspector.get_foreign_keys(tname, schema=schema)
         if autogen_context.run_name_filters(
@@ -1255,9 +1257,13 @@ def _compare_foreign_keys(
         )
     ]
 
-    backend_reflects_fk_options = bool(conn_fks and "options" in conn_fks[0])
+    backend_reflects_fk_options = bool(
+        conn_fks_list and "options" in conn_fks_list[0]
+    )
 
-    conn_fks = set(_make_foreign_key(const, conn_table) for const in conn_fks)
+    conn_fks = set(
+        _make_foreign_key(const, conn_table) for const in conn_fks_list
+    )
 
     # give the dialect a chance to correct the FKs to match more
     # closely
@@ -1265,24 +1271,24 @@ def _compare_foreign_keys(
         conn_fks, metadata_fks
     )
 
-    metadata_fks = set(
+    metadata_fks_sig = set(
         _fk_constraint_sig(fk, include_options=backend_reflects_fk_options)
         for fk in metadata_fks
     )
 
-    conn_fks = set(
+    conn_fks_sig = set(
         _fk_constraint_sig(fk, include_options=backend_reflects_fk_options)
         for fk in conn_fks
     )
 
-    conn_fks_by_sig = dict((c.sig, c) for c in conn_fks)
-    metadata_fks_by_sig = dict((c.sig, c) for c in metadata_fks)
+    conn_fks_by_sig = dict((c.sig, c) for c in conn_fks_sig)
+    metadata_fks_by_sig = dict((c.sig, c) for c in metadata_fks_sig)
 
     metadata_fks_by_name = dict(
-        (c.name, c) for c in metadata_fks if c.name is not None
+        (c.name, c) for c in metadata_fks_sig if c.name is not None
     )
     conn_fks_by_name = dict(
-        (c.name, c) for c in conn_fks if c.name is not None
+        (c.name, c) for c in conn_fks_sig if c.name is not None
     )
 
     def _add_fk(obj, compare_to):
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 77f0a86..9c992b4 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections import OrderedDict
 from io import StringIO
 import re
diff --git a/alembic/autogenerate/rewriter.py b/alembic/autogenerate/rewriter.py
index 0fdd398..79f665a 100644
--- a/alembic/autogenerate/rewriter.py
+++ b/alembic/autogenerate/rewriter.py
@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 from typing import Any
 from typing import Callable
 from typing import Iterator
 from typing import List
+from typing import Optional
 from typing import Type
 from typing import TYPE_CHECKING
 from typing import Union
@@ -49,12 +52,12 @@ class Rewriter:
 
     _traverse = util.Dispatcher()
 
-    _chained = None
+    _chained: Optional[Rewriter] = None
 
     def __init__(self) -> None:
         self.dispatch = util.Dispatcher()
 
-    def chain(self, other: "Rewriter") -> "Rewriter":
+    def chain(self, other: Rewriter) -> Rewriter:
         """Produce a "chain" of this :class:`.Rewriter` to another.
 
         This allows two rewriters to operate serially on a stream,