Make 'makemigrations --merge' update max_migrations.txt (#361)

Fixes #78.

CHANGELOG.rst

@@ -2,6 +2,10 @@
 Changelog
 =========
 
+* Make ``makemigrations --merge`` update ``max_migration.txt`` files as well.
+
+  Thanks to Gordon Wrigley for the report in `Issue #78 <https://github.com/adamchainz/django-linear-migrations/issues/78>`__.
+
 2.14.0 (2024-10-12)
 -------------------
 

src/django_linear_migrations/management/commands/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from django.db.migrations.writer import MigrationWriter
+
+
+@contextmanager
+def spy_on_migration_writers() -> Generator[dict[str, str]]:
+    written_migrations = {}
+
+    orig_as_string = MigrationWriter.as_string
+
+    def wrapped_as_string(self: MigrationWriter, *args: Any, **kwargs: Any) -> str:
+        written_migrations[self.migration.app_label] = self.migration.name
+        return orig_as_string(self, *args, **kwargs)
+
+    MigrationWriter.as_string = wrapped_as_string  # type: ignore [method-assign]
+    try:
+        yield written_migrations
+    finally:
+        MigrationWriter.as_string = orig_as_string  # type: ignore [method-assign]

src/django_linear_migrations/management/commands/makemigrations.py

@@ -1,54 +1,32 @@
 from __future__ import annotations
 
-import django
+from typing import Any
+
 from django.core.management.commands.makemigrations import Command as BaseCommand
-from django.db.migrations import Migration
 
 from django_linear_migrations.apps import MigrationDetails
 from django_linear_migrations.apps import first_party_app_configs
+from django_linear_migrations.management.commands import spy_on_migration_writers
 
 
 class Command(BaseCommand):
-    if django.VERSION >= (4, 2):
-
-        def write_migration_files(
-            self,
-            changes: dict[str, list[Migration]],
-            update_previous_migration_paths: dict[str, str] | None = None,
-        ) -> None:
-            # django-stubs awaiting new signature:
-            # https://github.com/typeddjango/django-stubs/pull/1609
-            super().write_migration_files(
-                changes,
-                update_previous_migration_paths,
-            )
-            _post_write_migration_files(self.dry_run, changes)
-
-    else:
-
-        def write_migration_files(  # type: ignore[misc,override]
-            self,
-            changes: dict[str, list[Migration]],
-        ) -> None:
-            super().write_migration_files(changes)
-            _post_write_migration_files(self.dry_run, changes)
-
-
-def _post_write_migration_files(
-    dry_run: bool, changes: dict[str, list[Migration]]
-) -> None:
-    if dry_run:
-        return
-
-    first_party_app_labels = {
-        app_config.label for app_config in first_party_app_configs()
-    }
-
-    for app_label, app_migrations in changes.items():
-        if app_label not in first_party_app_labels:
-            continue
-
-        # Reload required as we've generated changes
-        migration_details = MigrationDetails(app_label, do_reload=True)
-        max_migration_txt = migration_details.dir / "max_migration.txt"
-        max_migration_txt.write_text(f"{app_migrations[-1].name}\n")
+
+    def handle(self, *app_labels: Any, **options: Any) -> None:
+        with spy_on_migration_writers() as written_migrations:
+            super().handle(*app_labels, **options)
+
+        if options["dry_run"]:
+            return
+
+        first_party_app_labels = {
+            app_config.label for app_config in first_party_app_configs()
+        }
+
+        for app_label, migration_name in written_migrations.items():
+            if app_label not in first_party_app_labels:
+                continue
+
+            # Reload required in case of initial migration
+            migration_details = MigrationDetails(app_label, do_reload=True)
+            max_migration_txt = migration_details.dir / "max_migration.txt"
+            max_migration_txt.write_text(f"{migration_name}\n")

src/django_linear_migrations/management/commands/squashmigrations.py

@@ -2,38 +2,27 @@
 
 from typing import Any
 
-from django.core.management.commands import squashmigrations
 from django.core.management.commands.squashmigrations import Command as BaseCommand
-from django.db.migrations import Migration
-from django.db.migrations.writer import MigrationWriter
 
 from django_linear_migrations.apps import MigrationDetails
 from django_linear_migrations.apps import first_party_app_configs
+from django_linear_migrations.management.commands import spy_on_migration_writers
 
 
 class Command(BaseCommand):
     def handle(self, **options: Any) -> None:
-        # Temporarily wrap the call to MigrationWriter.__init__ to capture its first
-        # argument, the generated migration instance.
-        captured_migration = None
-
-        def wrapper(migration: Migration, *args: Any, **kwargs: Any) -> MigrationWriter:
-            nonlocal captured_migration
-            captured_migration = migration
-            return MigrationWriter(migration, *args, **kwargs)
+        with spy_on_migration_writers() as written_migrations:
+            super().handle(**options)
 
-        squashmigrations.MigrationWriter = wrapper  # type: ignore[attr-defined]
+        first_party_app_labels = {
+            app_config.label for app_config in first_party_app_configs()
+        }
 
-        try:
-            super().handle(**options)
-        finally:
-            squashmigrations.MigrationWriter = MigrationWriter  # type: ignore[attr-defined]
+        for app_label, migration_name in written_migrations.items():
+            if app_label not in first_party_app_labels:
+                continue
 
-        if captured_migration is not None and any(
-            captured_migration.app_label == app_config.label
-            for app_config in first_party_app_configs()
-        ):
             # A squash migration was generated, update max_migration.txt.
-            migration_details = MigrationDetails(captured_migration.app_label)
+            migration_details = MigrationDetails(app_label)
             max_migration_txt = migration_details.dir / "max_migration.txt"
-            max_migration_txt.write_text(f"{captured_migration.name}\n")
+            max_migration_txt.write_text(f"{migration_name}\n")

tests/test_makemigrations.py

@@ -86,3 +86,59 @@ def test_skips_creating_max_migration_txt_for_non_first_party_app(self):
         assert returncode == 0
         max_migration_txt = self.migrations_dir / "max_migration.txt"
         assert not max_migration_txt.exists()
+
+    def test_updates_for_a_merge(self):
+        (self.migrations_dir / "__init__.py").touch()
+        (self.migrations_dir / "0001_initial.py").write_text(
+            dedent(
+                """\
+            from django.db import migrations, models
+
+
+            class Migration(migrations.Migration):
+                initial = True
+                dependencies = []
+                operations = []
+            """
+            )
+        )
+        (self.migrations_dir / "0002_first_branch.py").write_text(
+            dedent(
+                """\
+            from django.db import migrations, models
+
+
+            class Migration(migrations.Migration):
+                dependencies = [
+                    ('testapp', '0001_initial'),
+                ]
+                operations = []
+            """
+            )
+        )
+        (self.migrations_dir / "0002_second_branch.py").write_text(
+            dedent(
+                """\
+            from django.db import migrations, models
+
+
+            class Migration(migrations.Migration):
+                dependencies = [
+                    ('testapp', '0001_initial'),
+                ]
+                operations = []
+            """
+            )
+        )
+        (self.migrations_dir / "max_migration.txt").write_text(
+            "0002_second_branch.py\n"
+        )
+
+        out, err, returncode = self.call_command("testapp", "--merge", "--no-input")
+
+        assert returncode == 0
+        max_migration_txt = self.migrations_dir / "max_migration.txt"
+        assert (
+            max_migration_txt.read_text()
+            == "0003_merge_0002_first_branch_0002_second_branch\n"
+        )