feat: Allow use class as argument to Register.migrate_data (#4)

src/nme/_class_register.py

@@ -3,6 +3,7 @@
 """
 import importlib
 import inspect
+import warnings
 from dataclasses import dataclass
 from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
@@ -35,6 +36,17 @@ def str_to_version(version: Union[str, Version]) -> Version:
     return parse_version(version) if isinstance(version, str) else version
 
 
+def _class_str_replace(func):
+    @wraps(func)
+    def _wrap(*args, **kwargs):
+        if "class_str" in kwargs:
+            warnings.warn("class_str argument is deprecated, use cls instead", FutureWarning, stacklevel=2)
+            kwargs["cls"] = kwargs.pop("class_str")
+        return func(*args, **kwargs)
+
+    return _wrap
+
+
 MigrationCallable = Callable[[Dict[str, Any]], Dict[str, Any]]
 MigrationInfo = Tuple[Version, MigrationCallable]
 """Type describing single migration entry. For given class Version number should be unique."""
@@ -151,8 +163,9 @@ def get_class(self, class_str: str) -> Type:
         self._register_missed(class_str=class_str)
         return self._data_dkt[class_str].type_
 
+    @_class_str_replace
     def migrate_data(
-        self, class_str: str, class_str_to_version_dkt: Dict[str, Union[str, Version]], data: Dict[str, Any]
+        self, cls: Union[str, Type], class_str_to_version_dkt: Dict[str, Union[str, Version]], data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """
         Apply migrations base on register state. Current implementation does not support multiple inheritance.
@@ -162,12 +175,15 @@ def migrate_data(
             If class is absent from this dict then assumed version is "0.0.0"
         :param data: dict of kwargs to constructor of class
         """
-        if self.use_parent_migrations(class_str):
-            super_klass = get_super_class(self.get_class(class_str))
+        if not isinstance(cls, str):
+            cls = class_to_str(cls)
+
+        if self.use_parent_migrations(cls):
+            super_klass = get_super_class(self.get_class(cls))
             if super_klass is not None:
                 data = self.migrate_data(class_to_str(super_klass), class_str_to_version_dkt, data)
-        version = str_to_version(class_str_to_version_dkt.get(class_str, "0.0.0"))
-        for version_, migration in self._data_dkt[class_str].migrations:
+        version = str_to_version(class_str_to_version_dkt.get(cls, "0.0.0"))
+        for version_, migration in self._data_dkt[cls].migrations:
             if version < version_:
                 data = migration(data)
         return data

src/tests/test_class_register.py

@@ -115,6 +115,33 @@ def __init__(self, field2, **kwargs):
     assert "field1" in migrated
 
 
+def test_migrate_by_class(clean_register):
+    @register_class(version="0.0.1", migrations=[("0.0.1", rename_key("field", "field1"))])
+    class MigrateClass:
+        def __init__(self, field1):
+            self.field1 = field1
+
+    migrated = REGISTER.migrate_data(class_to_str(MigrateClass), {}, {"field": 5})
+    migrated2 = REGISTER.migrate_data(MigrateClass, {}, {"field": 5})
+    assert migrated == {"field1": 5}
+    assert migrated == migrated2
+
+
+def test_migrate_keywords_args(clean_register):
+    @register_class(version="0.0.1", migrations=[("0.0.1", rename_key("field", "field1"))])
+    class MigrateClass:
+        def __init__(self, field1):
+            self.field1 = field1
+
+    migrated = REGISTER.migrate_data(cls=class_to_str(MigrateClass), class_str_to_version_dkt={}, data={"field": 5})
+    with pytest.warns(FutureWarning):
+        migrated2 = REGISTER.migrate_data(
+            class_str=class_to_str(MigrateClass), class_str_to_version_dkt={}, data={"field": 5}
+        )
+    assert migrated == {"field1": 5}
+    assert migrated == migrated2
+
+
 def test_wrong_version(clean_register):
     class SampleEnum(Enum):
         value1 = 1