I want to persist all attributes of an object which is an instance of a dataclass. Then I want to load back that object from the files that I persisted.
Here it is an example that fullfills the task:
JavaScript
x
30
30
1
from dataclasses import dataclass
2
import pickle
3
4
5
@dataclass
6
class Circle:
7
radius: float
8
centre: tuple
9
10
def save(self, path: str):
11
name = ".".join(("radius", "pkl"))
12
with open("/".join((path, name)), "wb") as f:
13
pickle.dump(self.radius, f)
14
name = ".".join(("centre", "pkl"))
15
with open("/".join((path, name)), "wb") as f:
16
pickle.dump(self.centre, f)
17
18
@classmethod
19
def load(cls, path):
20
my_model = {}
21
name = "radius"
22
file_name = ".".join((name, "pkl"))
23
with open("\".join((path, file_name)), "rb") as f:
24
my_model[name] = pickle.load(f)
25
name = "centre"
26
file_name = ".".join((name, "pkl"))
27
with open("\".join((path, file_name)), "rb") as f:
28
my_model[name] = pickle.load(f)
29
return cls(**my_model)
30
JavaScript
1
6
1
>>> c = Circle(2, (0, 0))
2
>>> c.save(r".Circle")
3
>>> c_loaded = Circle.load(r".Circle")
4
>>> c_loaded == c
5
True
6
As you can see I need to repeat the same code for every attribute, what is a better way to do it?
Advertisement
Answer
In the save method it use self.__dict__
. That contains all attribute names and values as a dictionary. Load is a classmethod so there is no __dict__
at that stage. However, cls.__annotations__
contains attribute names and types, still stored in a dictionary.
Here it is the end result:
JavaScript
1
23
23
1
from dataclasses import dataclass
2
import pickle
3
4
@dataclass
5
class Circle:
6
radius: float
7
centre: tuple
8
9
def save(self, path):
10
for name, attribute in self.__dict__.items():
11
name = ".".join((name, "pkl"))
12
with open("/".join((path, name)), "wb") as f:
13
pickle.dump(attribute, f)
14
15
@classmethod
16
def load(cls, path):
17
my_model = {}
18
for name in cls.__annotations__:
19
file_name = ".".join((name, "pkl"))
20
with open("/".join((path, file_name)), "rb") as f:
21
my_model[name] = pickle.load(f)
22
return cls(**my_model)
23