45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
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
|