interp-engine 0.0.24__py3-none-any.whl

neuron_explainer/fast_dataclasses/fast_dataclasses.py

@@ -0,0 +1,85 @@
+# Utilities for dataclasses that are very fast to serialize and deserialize, with limited data
+# validation. Fields must not be tuples, since they get serialized and then deserialized as lists.
+#
+# The unit tests for this library show how to use it.
+
+import json
+from dataclasses import dataclass, field, fields, is_dataclass
+from functools import partial
+from typing import Any, Union
+
+import orjson
+
+dataclasses_by_name = {}
+dataclasses_by_fieldnames = {}
+
+
+@dataclass
+class FastDataclass:
+    dataclass_name: str = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.dataclass_name = self.__class__.__name__
+
+
+def register_dataclass(cls):  # type: ignore
+    assert is_dataclass(cls), "Only dataclasses can be registered."
+    dataclasses_by_name[cls.__name__] = cls
+    name_set = frozenset(f.name for f in fields(cls) if f.name != "dataclass_name")
+    dataclasses_by_fieldnames[name_set] = cls
+    return cls
+
+
+def dumps(obj: Any) -> bytes:
+    return orjson.dumps(obj, option=orjson.OPT_SERIALIZE_NUMPY)
+
+
+def _object_hook(d: Any, backwards_compatible: bool = True) -> Any:
+    # If d is a list, recurse.
+    if isinstance(d, list):
+        return [_object_hook(x, backwards_compatible=backwards_compatible) for x in d]
+    # If d is not a dict, return it as is.
+    if not isinstance(d, dict):
+        return d
+    cls = None
+    if "dataclass_name" in d:
+        if d["dataclass_name"] in dataclasses_by_name:
+            cls = dataclasses_by_name[d["dataclass_name"]]
+        else:
+            assert backwards_compatible, (
+                f"Dataclass {d['dataclass_name']} not found, set backwards_compatible=True if you "
+                f"are okay with that."
+            )
+    # Load objects created without dataclass_name set.
+    else:
+        # Try our best to find a dataclass if backwards_compatible is True.
+        if backwards_compatible:
+            d_fields = frozenset(d.keys())
+            if d_fields in dataclasses_by_fieldnames:
+                cls = dataclasses_by_fieldnames[d_fields]
+            elif len(d_fields) > 0:
+                # Check if the fields are a subset of a dataclass (if the dataclass had extra fields
+                # added since the data was created). Note that this will fail if fields were removed
+                # from the dataclass.
+                for key, possible_cls in dataclasses_by_fieldnames.items():
+                    if d_fields.issubset(key):
+                        cls = possible_cls
+                        break
+                else:
+                    print(f"Could not find dataclass for {d_fields} {cls}")
+    new_d = {
+        k: _object_hook(v, backwards_compatible=backwards_compatible)
+        for k, v in d.items()
+        if k != "dataclass_name"
+    }
+    if cls is not None:
+        return cls(**new_d)
+    else:
+        return new_d
+
+
+def loads(s: Union[str, bytes], backwards_compatible: bool = True) -> Any:
+    return json.loads(
+        s,
+        object_hook=partial(_object_hook, backwards_compatible=backwards_compatible),
+    )

neuron_explainer/fast_dataclasses/test_fast_dataclasses.py

@@ -0,0 +1,83 @@
+from dataclasses import dataclass
+
+import pytest
+
+from .fast_dataclasses import FastDataclass, dumps, loads, register_dataclass
+
+
+# Inheritance is a bit tricky with our setup. dataclass_name must be set for instances of these
+# classes to serialize and deserialize correctly, but if it's given a default value, then subclasses
+# can't have any fields that don't have default values, because of how constructors are generated
+# for dataclasses (fields with no default value can't follow those with default values). To work
+# around this, we set dataclass_name in __post_init__ on the base class, which is called after the
+# constructor. The implementation does the right thing for both the base class and the subclass.
+@register_dataclass
+@dataclass
+class DataclassC(FastDataclass):
+    ints: list[int]
+
+
+@register_dataclass
+@dataclass
+class DataclassC_ext(DataclassC):
+    s: str
+
+
+@register_dataclass
+@dataclass
+class DataclassB(FastDataclass):
+    str_to_c: dict[str, DataclassC]
+    cs: list[DataclassC]
+
+
+@register_dataclass
+@dataclass
+class DataclassA(FastDataclass):
+    floats: list[float]
+    strings: list[str]
+    bs: list[DataclassB]
+
+
+@register_dataclass
+@dataclass
+class DataclassD(FastDataclass):
+    s1: str
+    s2: str = "default"
+
+
+def test_dataclasses() -> None:
+    a = DataclassA(
+        floats=[1.0, 2.0],
+        strings=["a", "b"],
+        bs=[
+            DataclassB(
+                str_to_c={"a": DataclassC(ints=[1, 2]), "b": DataclassC(ints=[3, 4])},
+                cs=[DataclassC(ints=[5, 6]), DataclassC_ext(ints=[7, 8], s="s")],
+            ),
+            DataclassB(
+                str_to_c={"c": DataclassC_ext(ints=[9, 10], s="t"), "d": DataclassC(ints=[11, 12])},
+                cs=[DataclassC(ints=[13, 14]), DataclassC(ints=[15, 16])],
+            ),
+        ],
+    )
+    assert loads(dumps(a)) == a
+
+
+def test_c_and_c_ext() -> None:
+    c_ext = DataclassC_ext(ints=[3, 4], s="s")
+    assert loads(dumps(c_ext)) == c_ext
+
+    c = DataclassC(ints=[1, 2])
+    assert loads(dumps(c)) == c
+
+
+def test_bad_serialized_data() -> None:
+    assert type(loads(dumps(DataclassC(ints=[3, 4])))) == DataclassC
+    assert type(loads('{"ints": [3, 4]}', backwards_compatible=False)) == dict
+    assert type(loads('{"ints": [3, 4], "dataclass_name": "DataclassC"}')) == DataclassC
+    with pytest.raises(TypeError):
+        loads('{"ints": [3, 4], "bogus_extra_field": "foo", "dataclass_name": "DataclassC"}')
+    with pytest.raises(TypeError):
+        loads('{"ints_field_is_missing": [3, 4], "dataclass_name": "DataclassC"}')
+    assert type(loads('{"s1": "test"}', backwards_compatible=False)) == dict
+    assert type(loads('{"s1": "test"}', backwards_compatible=True)) == DataclassD