Added typing to the default script mako templates.

Fixes: #1253 Change-Id: I00cae5daa8ebe9ba36bfac3af485a0e45a39f923
sqlalchemy · May 26, 2023 · 9705736 · 9705736
1 parent 78c23b6
commit 9705736
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 19 deletions.
diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
@@ -53,6 +53,7 @@
 
     from ..autogenerate.rewriter import Rewriter
     from ..runtime.migration import MigrationContext
+    from ..script.revision import _RevIdType
 
 
 class MigrateOperation:
@@ -2667,9 +2668,9 @@ def __init__(
         imports: Set[str] = set(),
         head: Optional[str] = None,
         splice: Optional[bool] = None,
-        branch_label: Optional[str] = None,
+        branch_label: Optional[_RevIdType] = None,
         version_path: Optional[str] = None,
-        depends_on: Optional[Union[str, Sequence[str]]] = None,
+        depends_on: Optional[_RevIdType] = None,
     ) -> None:
         self.rev_id = rev_id
         self.message = message

diff --git a/alembic/script/base.py b/alembic/script/base.py
@@ -26,23 +26,21 @@
 from ..util import not_none
 
 if TYPE_CHECKING:
+    from .revision import _RevIdType
+    from .revision import Revision
     from ..config import Config
     from ..config import MessagingOptions
     from ..runtime.migration import RevisionStep
     from ..runtime.migration import StampStep
-    from ..script.revision import Revision
 
 try:
     from dateutil import tz
 except ImportError:
     tz = None  # type: ignore[assignment]
 
-_RevIdType = Union[str, Sequence[str]]
-
 _sourceless_rev_file = re.compile(r"(?!\.\#|__init__)(.*\.py)(c|o)?$")
 _only_source_rev_file = re.compile(r"(?!\.\#|__init__)(.*\.py)$")
 _legacy_rev = re.compile(r"([a-f0-9]+)\.py$")
-_mod_def_re = re.compile(r"(upgrade|downgrade)_([a-z0-9]+)")
 _slug_re = re.compile(r"\w+")
 _default_file_template = "%(rev)s_%(slug)s"
 _split_on_space_comma = re.compile(r", *|(?: +)")
@@ -637,7 +635,7 @@ def generate_revision(
         head: Optional[str] = None,
         refresh: bool = False,
         splice: Optional[bool] = False,
-        branch_labels: Optional[str] = None,
+        branch_labels: Optional[_RevIdType] = None,
         version_path: Optional[str] = None,
         depends_on: Optional[_RevIdType] = None,
         **kw: Any,

diff --git a/alembic/templates/async/script.py.mako b/alembic/templates/async/script.py.mako
@@ -5,15 +5,17 @@ Revises: ${down_revision | comma,n}
 Create Date: ${create_date}
 
 """
+from typing import Sequence, Union
+
 from alembic import op
 import sqlalchemy as sa
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.
-revision = ${repr(up_revision)}
-down_revision = ${repr(down_revision)}
-branch_labels = ${repr(branch_labels)}
-depends_on = ${repr(depends_on)}
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
 
 
 def upgrade() -> None:

diff --git a/alembic/templates/generic/script.py.mako b/alembic/templates/generic/script.py.mako
@@ -5,15 +5,17 @@ Revises: ${down_revision | comma,n}
 Create Date: ${create_date}
 
 """
+from typing import Sequence, Union
+
 from alembic import op
 import sqlalchemy as sa
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.
-revision = ${repr(up_revision)}
-down_revision = ${repr(down_revision)}
-branch_labels = ${repr(branch_labels)}
-depends_on = ${repr(depends_on)}
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
 
 
 def upgrade() -> None:

diff --git a/alembic/templates/multidb/script.py.mako b/alembic/templates/multidb/script.py.mako
@@ -8,15 +8,17 @@ Revises: ${down_revision | comma,n}
 Create Date: ${create_date}
 
 """
+from typing import Sequence, Union
+
 from alembic import op
 import sqlalchemy as sa
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.
-revision = ${repr(up_revision)}
-down_revision = ${repr(down_revision)}
-branch_labels = ${repr(branch_labels)}
-depends_on = ${repr(depends_on)}
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
 
 
 def upgrade(engine_name: str) -> None:

diff --git a/docs/build/unreleased/1253.rst b/docs/build/unreleased/1253.rst
@@ -0,0 +1,5 @@
+.. change::
+    :tags: usecase, typing
+    :tickets: 1253
+
+    Added typing to the default script mako templates.