-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathutils.py
More file actions
36 lines (25 loc) · 910 Bytes
/
utils.py
File metadata and controls
36 lines (25 loc) · 910 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import numpy as np
import matplotlib.pyplot as plt
def figure_to_image(figure):
figure.canvas.draw()
img = np.frombuffer(figure.canvas.buffer_rgba(), dtype=np.uint8)
return img.reshape(figure.canvas.get_width_height()[::-1] + (4,))[..., :3]
def matches_post_pickle(figure):
import pickle
img_expected = figure_to_image(figure)
saved_fig = pickle.dumps(figure)
plt.close("all")
figure = pickle.loads(saved_fig)
img_result = figure_to_image(figure)
return np.all(img_expected == img_result)
def plotting_test(num_figs=1, *args, **kwargs):
def plotting_decorator(function):
def test_plotting():
plt.close("all")
res = function(
*(plt.figure(*args, **kwargs) for __ in range(num_figs))
)
plt.close("all")
return res
return test_plotting
return plotting_decorator