add save and load function for classes.

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