Enable providing migration dependencies as a list, while preserving b…

…ackward compatibility (#44)

* Enable providing migration dependencies as a list, while preserving backward compatibility

* Improving syntax

* Update annotations on migrator.after

Co-Authored-By: Łukasz Skarżyński <skarzynski_lukasz@protonmail.com>

Co-authored-by: Łukasz Skarżyński <skarzynski_lukasz@protonmail.com>

django_test_migrations/migrator.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from contextlib import contextmanager
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from django.core.management import call_command
 from django.db import DEFAULT_DB_ALIAS, connections
@@ -12,6 +12,7 @@
 # Regular or rollback migration: 0001 -> 0002, or 0002 -> 0001
 # Rollback migration to initial state: 0001 -> None
 _Migration = Tuple[str, Optional[str]]
+_MigrationSpec = Union[_Migration, List[_Migration]]
 
 
 @contextmanager
@@ -62,12 +63,14 @@ def __init__(
         self._database: str = database
         self._executor = MigrationExecutor(connections[self._database])
 
-    def before(self, migrate_from: _Migration) -> ProjectState:
+    def before(self, migrate_from: _MigrationSpec) -> ProjectState:
         """Reverse back to the original migration."""
+        if not isinstance(migrate_from, list):
+            migrate_from = [migrate_from]
         with _mute_migrate_signals():
-            return self._executor.migrate([migrate_from])
+            return self._executor.migrate(migrate_from)
 
-    def after(self, migrate_to: _Migration) -> ProjectState:
+    def after(self, migrate_to: _MigrationSpec) -> ProjectState:
         """Apply the next migration."""
         self._executor.loader.build_graph()  # reload.
         return self.before(migrate_to)

tests/test_migrator/test_migrator.py

@@ -16,3 +16,15 @@ def test_migrator(transactional_db):
     assert isinstance(old_state, ProjectState)
     assert isinstance(new_state, ProjectState)
     assert migrator.reset() is None
+
+
+@pytest.mark.django_db
+def test_migrator_list(transactional_db):
+    """We only need this test for coverage."""
+    migrator = Migrator()
+    old_state = migrator.before([('main_app', None)])
+    new_state = migrator.after([('main_app', '0001_initial')])
+
+    assert isinstance(old_state, ProjectState)
+    assert isinstance(new_state, ProjectState)
+    assert migrator.reset() is None