feat: recursive state casting (#12)

src/iokit/__init__.py

@@ -11,7 +11,7 @@
     "save_file",
     "save_temp",
 ]
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 
 from .extensions import Gzip, Json, Jsonl, Tar, Txt
 from .state import State, filter_states, find_state

src/iokit/state.py

@@ -4,10 +4,11 @@
     "find_state",
 ]
 
+from contextlib import suppress
 from datetime import datetime
 from fnmatch import fnmatch
 from io import BytesIO
-from typing import Any, Generator, Iterable
+from typing import Any, Generator, Iterable, Type
 
 from humanize import naturalsize
 from typing_extensions import Self
@@ -126,13 +127,17 @@ def __repr__(self) -> str:
         size = naturalsize(self.size, gnu=True)
         return f"{self.name} ({size})"
 
+    @classmethod
+    def by_suffix(cls, suffix: str) -> Type[Self]:
+        if suffix in cls._suffixes:
+            return cls
+        for klass in cls.__subclasses__():
+            with suppress(ValueError):
+                return klass.by_suffix(suffix)
+        raise ValueError(f"Unknown state suffix {suffix}")
+
     def cast(self) -> "State":
-        suffix = self.name.suffix
-        for klass in State.__subclasses__():
-            if suffix in getattr(klass, "_suffixes"):
-                break
-        else:
-            raise ValueError(f"Unknown state suffix {suffix}")
+        klass = self.by_suffix(self.name.suffix)
         state = klass.__new__(klass)
         setattr(state, "_data", self.data)
         setattr(state, "_name", self.name)

tests/test_state_inheritance.py

@@ -0,0 +1,14 @@
+from iokit import Json, State
+
+
+class MyJson(Json, suffix="myjson"):
+    pass
+
+
+def test_state_inheritance_json() -> None:
+    assert getattr(MyJson, "_suffix") == "myjson"
+    assert getattr(MyJson, "_suffixes") == ("myjson",)
+    myjson = MyJson({"a": 1}, name="test")
+    assert myjson.name == "test.myjson"
+    loaded = State(data=myjson.data, name="test.myjson")
+    assert loaded.load() == myjson.load() == {"a": 1}