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

kanta/kanta.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 from datetime import datetime
 from pathlib import Path
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 from typing import Any, Generic, TypeVar
 
 from kanta.kantaimpl import KantaImpl
@@ -50,7 +50,6 @@ class Kanta(Generic[T]):
         *,
         type: type[T] | None = None,
         migrations: ModuleType | str | None = None,
-        migration_ctx: Any | None = None,
         serializer: Serializer | None = None,
         flush_interval: float = 0.1,
     ):
@@ -61,7 +60,6 @@ class Kanta(Generic[T]):
             data: Caller-owned root msgspec.Struct state instance.
             type: Optional explicit root type. Defaults to ``type(data)``.
             migrations: Optional migrations module object or import path.
-            migration_ctx: Optional context object passed to migration functions.
             flush_interval: Background flush interval in seconds.
             serializer: Optional serializer implementation.
 
@@ -78,7 +76,6 @@ class Kanta(Generic[T]):
             data=data,
             type=data_type,
             migrations=migrations,
-            migration_ctx=migration_ctx,
             flush_interval=flush_interval,
             kanta=self,
         )
@@ -127,6 +124,15 @@ class Kanta(Generic[T]):
         """
         return self._impl.filename
 
+    @property
+    def ctx(self) -> SimpleNamespace:
+        """User-writable context namespace.
+
+        Migration functions receive the ``Kanta`` instance and can read or
+        mutate ``kanta.ctx`` during migrations.
+        """
+        return self._impl.ctx
+
     @property
     def mtime(self) -> datetime | None:
         """Last modification time carried forward from change records.

kanta/kantaimpl.py

@@ -7,6 +7,7 @@ import copy
 import importlib
 import logging
 from datetime import UTC, datetime
+from types import SimpleNamespace
 from typing import Any, Generic, TypeVar
 
 from kanta.callbacks import CallbackRegistry, InjectionContext
@@ -28,8 +29,8 @@ class KantaImpl(PersistenceMixin, Generic[T]):
         self.data_type = kwargs.pop("type")
         self.data: T = kwargs.pop("data")
         self.migrations = kwargs.pop("migrations", None)
-        self.migration_ctx = kwargs.pop("migration_ctx", None)
         self._kanta = kwargs.pop("kanta", None)
+        self.ctx = SimpleNamespace()
         super().__init__(**kwargs)
         self.migration_registry: MigrationRegistry | None = None
         if self.migrations is not None:
@@ -142,7 +143,7 @@ class KantaImpl(PersistenceMixin, Generic[T]):
 
             if self.migration_registry is not None:
                 rr.version = self.migration_registry.apply(
-                    rr.state, rr.version, self.migration_ctx
+                    rr.state, rr.version, self._kanta
                 )
 
             self.statedict = copy.deepcopy(rr.state)

kanta/migrate.py

@@ -7,24 +7,14 @@ or by prefix. Each runs exactly once based on the current version.
 from __future__ import annotations
 
 import importlib
+import inspect
 import logging
 from types import ModuleType
 from typing import Any
 
-import msgspec
-
 _logger = logging.getLogger(__name__)
 
 
-class MigrationCtx(msgspec.Struct, omit_defaults=True):
-    """Context passed to each migration function.
-
-    Subclass or replace this with your own context type.
-    """
-
-    pass
-
-
 class MigrationRegistry:
     """Registry of schema migration functions.
 
@@ -33,15 +23,20 @@ class MigrationRegistry:
         registry = MigrationRegistry()
 
         @registry.register
-        def migrate_v1(d: dict, ctx: MigrationCtx) -> None:
+        def migrate_v1(d: dict, kanta) -> None:
             d.setdefault("version", 1)
+            kanta.ctx.note = "migrated"
 
-        new_version = registry.apply(state, current_version=0)
+        @registry.register
+        def migrate_v2(d: dict) -> None:
+            d.setdefault("version", 2)
+
+        new_version = registry.apply(state, current_version=0, kanta=kanta)
 
     Or load from a module::
 
         registry = MigrationRegistry.from_module("myapp.migrations")
-        new_version = registry.apply(state, current_version=0)
+        new_version = registry.apply(state, current_version=0, kanta=kanta)
     """
 
     def __init__(self) -> None:
@@ -89,11 +84,21 @@ class MigrationRegistry:
         """Current schema version (= highest discovered migration, or 0)."""
         return max(self._migrations.keys(), default=0)
 
+    @staticmethod
+    def _call_migration(fn: Any, data_dict: dict[str, Any], kanta: Any) -> None:
+        """Call *fn* with the data dict and, if accepted, the Kanta instance."""
+        try:
+            inspect.signature(fn).bind(data_dict, kanta)
+        except TypeError:
+            fn(data_dict)
+        else:
+            fn(data_dict, kanta)
+
     def apply(
         self,
         data_dict: dict[str, Any],
         current_version: int,
-        ctx: MigrationCtx | None = None,
+        kanta: Any,
         *,
         silent: bool = False,
     ) -> int:
@@ -109,7 +114,7 @@ class MigrationRegistry:
                     f"Missing migration step migrate_v{next_version} "
                     f"(highest discovered is v{self.dbver})"
                 )
-            fn(data_dict, ctx or MigrationCtx())
+            self._call_migration(fn, data_dict, kanta)
             current_version = next_version
             if not silent:
                 desc = (fn.__doc__ or fn.__name__).split("\n")[0].rstrip(".")

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)