From 3e48ed0d4287fd89e96dce8ed1dd46b4c25bf53d Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 17 Oct 2023 19:18:20 -0400
Subject: [PATCH] match process_revision_directives typing to API

Improved typing in the
:paramref:`.EnvironmentContext.configure.process_revision_directives`
callable to better indicate that the passed-in type is
:class:`.MigrationScript`, not the :class:`.MigrationOperation` base class,
and added typing to the example at :ref:`cookbook_no_empty_migrations` to
illustrate.

Change-Id: Ibfb7a57a081818c290cf0964d12a72b85c2c1983
Fixes: #1325
---
 alembic/autogenerate/api.py    |  3 +++
 alembic/context.pyi            |  4 ++--
 alembic/operations/__init__.py |  2 ++
 alembic/runtime/environment.py |  4 ++--
 docs/build/cookbook.rst        | 14 +++++++++++++-
 docs/build/unreleased/1325.rst | 10 ++++++++++
 6 files changed, 32 insertions(+), 5 deletions(-)
 create mode 100644 docs/build/unreleased/1325.rst

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 9b7a97d4..13d025b2 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -5,6 +5,7 @@ from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import Iterator
+from typing import List
 from typing import Optional
 from typing import Sequence
 from typing import Set
@@ -508,6 +509,8 @@ class RevisionContext:
     """Maintains configuration and state that's specific to a revision
     file generation operation."""
 
+    generated_revisions: List[MigrationScript]
+
     def __init__(
         self,
         config: Config,
diff --git a/alembic/context.pyi b/alembic/context.pyi
index 5c093012..85e0cf75 100644
--- a/alembic/context.pyi
+++ b/alembic/context.pyi
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
 
     from .autogenerate.api import AutogenContext
     from .config import Config
-    from .operations.ops import MigrateOperation
+    from .operations.ops import MigrationScript
     from .runtime.migration import _ProxyTransaction
     from .runtime.migration import MigrationContext
     from .runtime.migration import MigrationInfo
@@ -143,7 +143,7 @@ def configure(
     include_schemas: bool = False,
     process_revision_directives: Optional[
         Callable[
-            [MigrationContext, Tuple[str, str], List[MigrateOperation]], None
+            [MigrationContext, Tuple[str, str], List[MigrationScript]], None
         ]
     ] = None,
     compare_type: Union[
diff --git a/alembic/operations/__init__.py b/alembic/operations/__init__.py
index 9de1918c..26197cbe 100644
--- a/alembic/operations/__init__.py
+++ b/alembic/operations/__init__.py
@@ -3,6 +3,7 @@ from .base import AbstractOperations
 from .base import BatchOperations
 from .base import Operations
 from .ops import MigrateOperation
+from .ops import MigrationScript
 
 
 __all__ = [
@@ -10,4 +11,5 @@ __all__ = [
     "Operations",
     "BatchOperations",
     "MigrateOperation",
+    "MigrationScript",
 ]
diff --git a/alembic/runtime/environment.py b/alembic/runtime/environment.py
index 18840470..a1c0e1b0 100644
--- a/alembic/runtime/environment.py
+++ b/alembic/runtime/environment.py
@@ -36,13 +36,13 @@ if TYPE_CHECKING:
     from ..autogenerate.api import AutogenContext
     from ..config import Config
     from ..ddl import DefaultImpl
-    from ..operations.ops import MigrateOperation
+    from ..operations.ops import MigrationScript
     from ..script.base import ScriptDirectory
 
 _RevNumber = Optional[Union[str, Tuple[str, ...]]]
 
 ProcessRevisionDirectiveFn = Callable[
-    [MigrationContext, Tuple[str, str], List["MigrateOperation"]], None
+    [MigrationContext, Tuple[str, str], List["MigrationScript"]], None
 ]
 
 RenderItemFn = Callable[
diff --git a/docs/build/cookbook.rst b/docs/build/cookbook.rst
index 510699a3..f5cf6cbb 100644
--- a/docs/build/cookbook.rst
+++ b/docs/build/cookbook.rst
@@ -900,14 +900,26 @@ hook in :meth:`.MigrationContext.configure` which removes the
 single :class:`.MigrationScript` directive if it is empty of
 any operations::
 
+    # for typing purposes
+    from alembic.environment import MigrationContext
+
+    # this typing-only import requires alembic 1.12.1 or above
+    from alembic.operations import MigrationScript
+
 
     def run_migrations_online():
 
         # ...
 
-        def process_revision_directives(context, revision, directives):
+        def process_revision_directives(
+            context: MigrationContext,
+            revision: tuple[str, str],
+            directives: list[MigrationScript],
+        ):
+            assert config.cmd_opts is not None
             if config.cmd_opts.autogenerate:
                 script = directives[0]
+                assert script.upgrade_ops is not None
                 if script.upgrade_ops.is_empty():
                     directives[:] = []
 
diff --git a/docs/build/unreleased/1325.rst b/docs/build/unreleased/1325.rst
new file mode 100644
index 00000000..1b4ccfaa
--- /dev/null
+++ b/docs/build/unreleased/1325.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 1325
+
+    Improved typing in the
+    :paramref:`.EnvironmentContext.configure.process_revision_directives`
+    callable to better indicate that the passed-in type is
+    :class:`.MigrationScript`, not the :class:`.MigrationOperation` base class,
+    and added typing to the example at :ref:`cookbook_no_empty_migrations` to
+    illustrate.
-- 
2.47.2