forked from t-makaro/animatplot
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtools.py
More file actions
91 lines (72 loc) · 3.24 KB
/
tools.py
File metadata and controls
91 lines (72 loc) · 3.24 KB
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import os
from os.path import join, split
import functools
import pytest
import matplotlib.pyplot as plt
from matplotlib.animation import FileMovieWriter
from matplotlib.testing.compare import compare_images
from matplotlib.testing.decorators import remove_ticks_and_titles
from matplotlib.testing.exceptions import ImageComparisonFailure
class BunchOFiles(FileMovieWriter):
"""
Exports an animation to a series of images.
called like:
``anim.save('name.format', writer=BunchOFiles())``
which produces a series of files named ``name{num}.format``
"""
supported_formats = ['png', 'jpeg', 'bmp', 'svg', 'pdf']
def __init__(self, *args, extra_args=None, **kwargs):
# extra_args aren't used but we need to stop None from being passed
super().__init__(*args, extra_args=(), **kwargs)
def setup(self, fig, dpi, frame_prefix):
super().setup(fig, dpi, frame_prefix, clear_temp=False)
self.fname_format_str = '%s%%d.%s'
self.temp_prefix, self.frame_format = self.outfile.split('.')
def grab_frame(self, **savefig_kwargs):
'''
Grab the image information from the figure and save as a movie frame.
All keyword arguments in savefig_kwargs are passed on to the 'savefig'
command that saves the figure.
'''
# Tell the figure to save its data to the sink, using the
# frame format and dpi.
with self._frame_sink() as myframesink:
self.fig.savefig(myframesink, format=self.frame_format,
dpi=self.dpi, **savefig_kwargs)
def finish(self):
self._frame_sink().close()
def _compare_animation(anim, expected, format_, nframes, tol):
# generate images from the animation
base_dir, filename = split(join('tests', 'baseline_images', expected))
out_dir = split(join('tests', 'output_images', expected))[0]
if not os.path.exists(out_dir):
os.makedirs(out_dir)
anim.save(os.path.join(out_dir, (filename+format_)), writer=BunchOFiles())
for i in range(nframes):
image_name = '%s%d%s' % (filename, i, format_)
expected_name = os.path.join(base_dir, image_name)
actual_name = os.path.join(out_dir, image_name)
err = compare_images(expected_name, actual_name, tol,
in_decorator=True)
if not os.path.exists(expected_name):
raise ImageComparisonFailure('image does not exist: %s' % expected_name)
if err:
for key in ["actual", "expected"]:
err[key] = os.path.relpath(err[key])
raise ImageComparisonFailure(
'images not close (RMS %(rms).3f):\n\t%(actual)s\n\t%(expected)s ' % err)
def animation_compare(baseline_images, nframes, fmt='.png', tol=1e-3, remove_text=True):
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
anim = func(*args, **kwargs)
if remove_text:
fignum = plt.get_fignums()[0]
fig = plt.figure(fignum)
remove_ticks_and_titles(fig)
try:
_compare_animation(anim, baseline_images, fmt, nframes, tol)
finally:
plt.close('all')
return wrapper
return decorator