Added typing to the default script mako templates.

Fixes: #1253
Change-Id: I00cae5daa8ebe9ba36bfac3af485a0e45a39f923

diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index c65ed16c76df6d9d4791c3cfb8fc306ffd4664a4..ae929fae53d51b3db878537e11ceff1f4864ca60 100644 (file)
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -53,6 +53,7 @@ if TYPE_CHECKING:
 
     from ..autogenerate.rewriter import Rewriter
     from ..runtime.migration import MigrationContext
+    from ..script.revision import _RevIdType
 
 
 class MigrateOperation:
@@ -2667,9 +2668,9 @@ class MigrationScript(MigrateOperation):
         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
index 244086589f930086fa375593759a2fc6ed6f27ca..cf929c978d2d6da06ad59e25310c8fbb6a192943 100644 (file)
--- a/alembic/script/base.py
+++ b/alembic/script/base.py
@@ -26,23 +26,21 @@ from ..runtime import migration
 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 @@ class ScriptDirectory:
         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
index 55df2863d206fa1678abb4c92e90c45d3f85c114..fbc4b07dcef98b20c6f96b642097f35e8433258e 100644 (file)
--- 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
index 55df2863d206fa1678abb4c92e90c45d3f85c114..fbc4b07dcef98b20c6f96b642097f35e8433258e 100644 (file)
--- 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
index 946ab6a9740355de8808a18ad1df37a17bc929e3..6108b8a0dc2a1bef8f8b25fda188ab61ef9bdaf9 100644 (file)
--- 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
new file mode 100644 (file)
index 0000000..592cc60
--- /dev/null
+++ b/docs/build/unreleased/1253.rst
@@ -0,0 +1,5 @@
+.. change::
+    :tags: usecase, typing
+    :tickets: 1253
+
+    Added typing to the default script mako templates.