Refactor Migrations internals, add caching to avoid reloading per each Kanta instance. Naming changed from MigrationRegistry to Migration, module from migrate to kanta.migrations.

kanta/kantaimpl.py

@@ -12,7 +12,7 @@ from typing import Any, Generic, TypeVar
 
 from kanta.callbacks import CallbackRegistry, InjectionContext
 from kanta.exceptions import DatabaseError, DataIntegrityError, ReplayError
-from kanta.migrate import MigrationRegistry
+from kanta.migrations import Migrations
 from kanta.persistence import PersistenceMixin
 from kanta.serialization import restore_data_in_place, struct_to_dict
 from kanta.serialization.base import replay
@@ -28,18 +28,18 @@ class KantaImpl(PersistenceMixin, Generic[T]):
     def __init__(self, **kwargs: Any):
         self.data_type = kwargs.pop("type")
         self.data: T = kwargs.pop("data")
-        self.migrations = kwargs.pop("migrations", None)
         self._kanta = kwargs.pop("kanta", None)
+        migrations = kwargs.pop("migrations", None)
         self.ctx = SimpleNamespace()
         super().__init__(**kwargs)
-        self.migration_registry: MigrationRegistry | None = None
+        self.migrations: Migrations | None = None
-        if self.migrations is not None:
+        if migrations is not None:
             module = (
-                importlib.import_module(self.migrations)
+                importlib.import_module(migrations)
-                if isinstance(self.migrations, str)
+                if isinstance(migrations, str)
-                else self.migrations
+                else migrations
             )
-            self.migration_registry = MigrationRegistry.from_module(module)
+            self.migrations = Migrations.from_module(module)
 
         self.in_transaction = False
         self.transaction_snapshot: dict[str, Any] | None = None
@@ -55,9 +55,7 @@ class KantaImpl(PersistenceMixin, Generic[T]):
         )
 
         self.statedict = struct_to_dict(self.data, serializer=self.serializer)
-        self.version = (
+        self.version = self.migrations.dbver if self.migrations is not None else 0
-            self.migration_registry.dbver if self.migration_registry is not None else 0
-        )
 
     def add_bootstrap(
         self,
@@ -141,10 +139,8 @@ class KantaImpl(PersistenceMixin, Generic[T]):
                     cause_type=type(e).__name__,
                 ) from e
 
-            if self.migration_registry is not None:
+            if self.migrations is not None:
-                rr.version = self.migration_registry.apply(
+                rr.version = self.migrations.apply(rr.state, rr.version, self._kanta)
-                    rr.state, rr.version, self._kanta
-                )
 
             self.statedict = copy.deepcopy(rr.state)
             self.data = restore_data_in_place(

kanta/migrations.py

@@ -14,29 +14,33 @@ from typing import Any
 
 _logger = logging.getLogger(__name__)
 
+# Cache registries by imported module object so that many Kanta instances using
+# the same migrations module do not re-scan it each time.
+_module_registry_cache: dict[ModuleType, Migrations] = {}
 
-class MigrationRegistry:
+
+class Migrations:
     """Registry of schema migration functions.
 
     Usage::
 
-        registry = MigrationRegistry()
+        migrations = Migrations()
 
-        @registry.register
+        @migrations.register
         def migrate_v1(d: dict, kanta) -> None:
             d.setdefault("version", 1)
             kanta.ctx.note = "migrated"
 
-        @registry.register
+        @migrations.register
         def migrate_v2(d: dict) -> None:
             d.setdefault("version", 2)
 
-        new_version = registry.apply(state, current_version=0, kanta=kanta)
+        new_version = migrations.apply(state, current_version=0, kanta=kanta)
 
     Or load from a module::
 
-        registry = MigrationRegistry.from_module("myapp.migrations")
+        migrations = Migrations.from_module("myapp.migrations")
-        new_version = registry.apply(state, current_version=0, kanta=kanta)
+        new_version = migrations.apply(state, current_version=0, kanta=kanta)
     """
 
     def __init__(self) -> None:
@@ -59,24 +63,30 @@ class MigrationRegistry:
         return fn
 
     @classmethod
-    def from_module(cls, module: str | ModuleType) -> MigrationRegistry:
+    def from_module(cls, module: str | ModuleType) -> Migrations:
-        """Create a registry by scanning a module for ``migrate_vN`` functions.
+        """Create or retrieve a cached registry by scanning a module.
 
         Args:
             module: A module name (string) or an imported module object.
         """
-        reg = cls()
         if isinstance(module, str):
             mod = importlib.import_module(module)
         else:
             mod = module
 
+        try:
+            return _module_registry_cache[mod]
+        except KeyError:
+            pass
+
+        reg = cls()
         for name in dir(mod):
             if name.startswith("migrate_v"):
                 fn = getattr(mod, name)
                 if callable(fn):
                     version = reg._migration_version(fn)
                     reg._migrations[version] = fn
+        _module_registry_cache[mod] = reg
         return reg
 
     @property

tests/test_migrations.py

@@ -1,6 +1,6 @@
 from types import ModuleType, SimpleNamespace
 
-from kanta.migrate import MigrationRegistry
+from kanta.migrations import Migrations
 
 
 class _DummyKanta:
@@ -9,7 +9,7 @@ class _DummyKanta:
 
 
 def test_register_and_apply():
-    reg = MigrationRegistry()
+    reg = Migrations()
     kanta = _DummyKanta()
 
     @reg.register
@@ -27,7 +27,7 @@ def test_register_and_apply():
 
 
 def test_no_migrations_needed():
-    reg = MigrationRegistry()
+    reg = Migrations()
     kanta = _DummyKanta()
 
     @reg.register
@@ -52,7 +52,7 @@ def test_from_module():
     mod.__dict__["migrate_v1"] = migrate_v1
     mod.__dict__["migrate_v2"] = migrate_v2
 
-    reg = MigrationRegistry.from_module(mod)
+    reg = Migrations.from_module(mod)
     assert reg.dbver == 2
 
     state = {}
@@ -62,7 +62,7 @@ def test_from_module():
 
 
 def test_migrations_can_use_kanta_ctx():
-    reg = MigrationRegistry()
+    reg = Migrations()
     kanta = _DummyKanta()
 
     @reg.register
@@ -78,7 +78,7 @@ def test_migrations_can_use_kanta_ctx():
 
 
 def test_migration_can_omit_kanta_argument():
-    reg = MigrationRegistry()
+    reg = Migrations()
     kanta = _DummyKanta()
 
     @reg.register