Add kanta.ctx SimpleNamespace for user variables. This is passed to migration functions if they take a second argument. Remove the old MigrationCtx system.

tests/test_kanta_integration.py

@@ -404,7 +404,7 @@ async def test_migrations_from_module(tmp_path, format_config):
 
     mod = type(sys)("test_migrations")
 
-    def migrate_v1(d, ctx):
+    def migrate_v1(d, kanta):
         d["version"] = 1
 
     mod.__dict__["migrate_v1"] = migrate_v1
@@ -514,7 +514,7 @@ async def test_migrations_from_module_path(tmp_path, format_config):
     module_name = "test_migrations_path"
     mod = type(sys)(module_name)
 
-    def migrate_v1(d, ctx):
+    def migrate_v1(d, kanta):
         d["counter"] = 2
 
     mod.__dict__["migrate_v1"] = migrate_v1

tests/test_migrations.py

@@ -1,44 +1,52 @@
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 
 from kanta.migrate import MigrationRegistry
 
 
+class _DummyKanta:
+    def __init__(self):
+        self.ctx = SimpleNamespace()
+
+
 def test_register_and_apply():
     reg = MigrationRegistry()
+    kanta = _DummyKanta()
 
     @reg.register
-    def migrate_v1(d, ctx):
+    def migrate_v1(d, kanta):
         d["version"] = 1
 
     @reg.register
-    def migrate_v2(d, ctx):
+    def migrate_v2(d, kanta):
         d["version"] = 2
 
     state = {}
-    new_ver = reg.apply(state, current_version=0, silent=True)
+    new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True)
     assert new_ver == 2
     assert state["version"] == 2
 
 
 def test_no_migrations_needed():
     reg = MigrationRegistry()
+    kanta = _DummyKanta()
 
     @reg.register
-    def migrate_v1(d, ctx):
+    def migrate_v1(d, kanta):
         d["x"] = 1
 
     state = {"x": 1}
-    new_ver = reg.apply(state, current_version=1, silent=True)
+    new_ver = reg.apply(state, current_version=1, kanta=kanta, silent=True)
     assert new_ver == 1
 
 
 def test_from_module():
     mod = ModuleType("fake_migrations")
+    kanta = _DummyKanta()
 
-    def migrate_v1(d, ctx):
+    def migrate_v1(d, kanta):
         d["v"] = 1
 
-    def migrate_v2(d, ctx):
+    def migrate_v2(d, kanta):
         d["v"] = 2
 
     mod.__dict__["migrate_v1"] = migrate_v1
@@ -48,6 +56,36 @@ def test_from_module():
     assert reg.dbver == 2
 
     state = {}
-    new_ver = reg.apply(state, current_version=0, silent=True)
+    new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True)
     assert new_ver == 2
     assert state["v"] == 2
+
+
+def test_migrations_can_use_kanta_ctx():
+    reg = MigrationRegistry()
+    kanta = _DummyKanta()
+
+    @reg.register
+    def migrate_v1(d, kanta):
+        kanta.ctx.source = "migration"
+        d["source"] = kanta.ctx.source
+
+    state = {}
+    new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True)
+    assert new_ver == 1
+    assert state["source"] == "migration"
+    assert kanta.ctx.source == "migration"
+
+
+def test_migration_can_omit_kanta_argument():
+    reg = MigrationRegistry()
+    kanta = _DummyKanta()
+
+    @reg.register
+    def migrate_v1(d):
+        d["x"] = 1
+
+    state = {}
+    new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True)
+    assert new_ver == 1
+    assert state["x"] == 1

tests/test_readonly.py

@@ -116,7 +116,7 @@ async def test_readonly_runs_migrations(tmp_path, format_config):
         format_config,
     )
 
-    def migrate_v1(data, ctx):
+    def migrate_v1(data, kanta):
         data.setdefault("enabled", True)
 
     migrations = make_migrations_module("readonly_migrations", "migrate_v1", migrate_v1)