fix dataset by fixing getstate

misko · Dec 29, 2023 · ce9e4aa · ce9e4aa
1 parent 0cf4778
commit ce9e4aa
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/spf/utils.py b/spf/utils.py
@@ -0,0 +1,8 @@
+class dotdict(dict):
+    __getattr__ = dict.get
+
+    def __getstate__(self):
+        return vars(self)
+
+    def __setstate__(self, state):
+        vars(self).update(state)
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -5,13 +5,7 @@
 
 from spf.dataset.spf_generate import generate_session_and_dump
 
-
-class dotdict(dict):
-    """dot.notation access to dictionary attributes"""
-
-    __getattr__ = dict.get
-    __setattr__ = dict.__setitem__
-    __delattr__ = dict.__delitem__
+from spf.utils import dotdict
 
 
 @pytest.fixture
@@ -52,7 +46,7 @@ def test_data_generation(default_args):
         args = default_args
         args.output = tmp
         dump(
-            {k: v for k, v in args.items()},
+            args,
             "/".join([args.output, "args.pkl"]),
             compression="lzma",
         )
@@ -68,7 +62,7 @@ def test_live_data_generation(default_args):
         args = default_args
         args.output = tmp
         dump(
-            {k: v for k, v in args.items()},
+            args,
             "/".join([args.output, "args.pkl"]),
             compression="lzma",
         )