From: Federico Caselli <cfederico87@gmail.com>
Date: Fri, 26 May 2023 20:01:48 +0000 (+0200)
Subject: Added typing to the default script mako templates.
X-Git-Tag: rel_1_11_2~5
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=97057360e5780f467a71f716cd2c91d698d77448;p=thirdparty%2Fsqlalchemy%2Falembic.git

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 c65ed16c..ae929fae 100644
--- 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 24408658..cf929c97 100644
--- 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 55df2863..fbc4b07d 100644
--- 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 55df2863..fbc4b07d 100644
--- 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 946ab6a9..6108b8a0 100644
--- 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
index 00000000..592cc602
--- /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.