From c998ece8feac9136f24e671ce311555e08fca64a Mon Sep 17 00:00:00 2001
From: Stephen Brannen <stephen@trykintsugi.com>
Date: Mon, 2 Mar 2026 16:05:36 -0600
Subject: [PATCH 1/2] fix: handle annotated assignments (AnnAssign) in
 migration files

Newer Alembic versions generate `revision: str = '...'` instead of
`revision = '...'`. The AST parser only handled Assign nodes, silently
skipping every migration that used the annotated form.

Closes #4
---
 squawk_alembic/hook.py      | 16 +++++++----
 tests/test_revision_info.py | 57 +++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/squawk_alembic/hook.py b/squawk_alembic/hook.py
index 58a0595..b6508cd 100644
--- a/squawk_alembic/hook.py
+++ b/squawk_alembic/hook.py
@@ -6,7 +6,7 @@
 import subprocess
 import sys
 import tempfile
-from ast import Assign, Constant, Name, Tuple, iter_child_nodes, parse
+from ast import AnnAssign, Assign, Constant, Name, Tuple, iter_child_nodes, parse
 from configparser import ConfigParser, NoOptionError, NoSectionError
 from pathlib import Path
 
@@ -57,12 +57,16 @@ def extract_revision_info(filepath):
     down_revision = None
 
     for node in iter_child_nodes(tree):
-        if not isinstance(node, Assign):
+        if isinstance(node, AnnAssign):
+            if not isinstance(node.target, Name) or node.value is None:
+                continue
+            name = node.target.id
+        elif isinstance(node, Assign):
+            if len(node.targets) != 1 or not isinstance(node.targets[0], Name):
+                continue
+            name = node.targets[0].id
+        else:
             continue
-        if len(node.targets) != 1 or not isinstance(node.targets[0], Name):
-            continue
-
-        name = node.targets[0].id
         if name == "revision":
             if isinstance(node.value, Constant) and isinstance(node.value.value, str):
                 revision = node.value.value
diff --git a/tests/test_revision_info.py b/tests/test_revision_info.py
index d808196..166fd6c 100644
--- a/tests/test_revision_info.py
+++ b/tests/test_revision_info.py
@@ -85,3 +85,60 @@ def upgrade():
             pass
     """)
     assert extract_revision_info(path) is None
+
+
+def test_annotated_assignment(migration_file):
+    path = migration_file("""
+        from typing import Sequence, Union
+
+        revision: str = 'abc123'
+        down_revision: Union[str, None] = 'def456'
+        branch_labels: Union[str, Sequence[str], None] = None
+        depends_on: Union[str, Sequence[str], None] = None
+
+        def upgrade():
+            pass
+    """)
+    info = extract_revision_info(path)
+    assert info is not None
+    assert info.revision == "abc123"
+    assert info.down_revision == "def456"
+    assert info.is_merge is False
+
+
+def test_annotated_first_migration_down_revision_is_none(migration_file):
+    path = migration_file("""
+        from typing import Sequence, Union
+
+        revision: str = 'abc123'
+        down_revision: Union[str, None] = None
+        branch_labels: Union[str, Sequence[str], None] = None
+        depends_on: Union[str, Sequence[str], None] = None
+
+        def upgrade():
+            pass
+    """)
+    info = extract_revision_info(path)
+    assert info is not None
+    assert info.revision == "abc123"
+    assert info.down_revision is None
+    assert info.is_merge is False
+
+
+def test_annotated_merge_migration(migration_file):
+    path = migration_file("""
+        from typing import Sequence, Union
+
+        revision: str = 'merge001'
+        down_revision: Union[str, None] = ('abc123', 'def456')
+        branch_labels: Union[str, Sequence[str], None] = None
+        depends_on: Union[str, Sequence[str], None] = None
+
+        def upgrade():
+            pass
+    """)
+    info = extract_revision_info(path)
+    assert info is not None
+    assert info.revision == "merge001"
+    assert info.down_revision == ("abc123", "def456")
+    assert info.is_merge is True

From e80ff29c63158eb709bdd0c5da6569e1a1b470c3 Mon Sep 17 00:00:00 2001
From: Stephen Brannen <stephen@trykintsugi.com>
Date: Mon, 2 Mar 2026 16:06:19 -0600
Subject: [PATCH 2/2] chore: bump version to 0.3.2

---
 pyproject.toml             | 2 +-
 squawk_alembic/__init__.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8e3edab..cf5806f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "squawk-alembic"
-version = "0.3.1"
+version = "0.3.2"
 description = "Pre-commit hook to lint Alembic migration SQL with squawk"
 packages = [{include = "squawk_alembic"}]
 readme = "README.md"
diff --git a/squawk_alembic/__init__.py b/squawk_alembic/__init__.py
index 260c070..f9aa3e1 100644
--- a/squawk_alembic/__init__.py
+++ b/squawk_alembic/__init__.py
@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"