add save and load function for classes.

tensorneat/utils/__init__.py

@@ -3,3 +3,4 @@ from .aggregation import Agg, agg_func, AGG_ALL
 from .tools import *
 from .graph import *
 from .state import State
+from .stateful_class import StatefulBaseClass

tensorneat/utils/state.py

@@ -30,6 +30,12 @@ class State:
     def __repr__(self):
         return f"State ({self.state_dict})"
 
+    def __getstate__(self):
+        return self.state_dict.copy()
+
+    def __setstate__(self, state):
+        self.__dict__["state_dict"] = state
+
     def tree_flatten(self):
         children = list(self.state_dict.values())
         aux_data = list(self.state_dict.keys())

tensorneat/utils/stateful_class.py

@@ -0,0 +1,44 @@
+from typing import Optional
+from . import State
+import pickle
+import datetime
+import warnings
+
+
+class StatefulBaseClass:
+    def setup(self, state=State()):
+        return state
+
+    def save(self, state: Optional[State] = None, path: Optional[str] = None):
+        if path is None:
+            time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+            path = f"./{self.__class__.__name__} {time}.pkl"
+        if state is not None:
+            self.__dict__["aux_for_state"] = state
+        with open(path, "wb") as f:
+            pickle.dump(self, f)
+
+    @classmethod
+    def load(cls, path: str, with_state: bool = False, warning: bool = True):
+        with open(path, "rb") as f:
+            obj = pickle.load(f)
+        if with_state:
+            if "aux_for_state" not in obj.__dict__:
+                if warning:
+                    warnings.warn(
+                        "This object does not have state to load, return empty state",
+                        category=UserWarning,
+                    )
+                return obj, State()
+            state = obj.__dict__["aux_for_state"]
+            del obj.__dict__["aux_for_state"]
+            return obj, state
+        else:
+            if "aux_for_state" in obj.__dict__:
+                if warning:
+                    warnings.warn(
+                        "This object state to load, ignore it",
+                        category=UserWarning,
+                    )
+                del obj.__dict__["aux_for_state"]
+            return obj