34
34
import os
35
35
import sys
36
36
import shutil
37
+ import inspect
37
38
import tempfile
38
39
import warnings
39
40
from distutils .version import LooseVersion
42
43
43
44
if sys .version_info [0 ] == 2 :
44
45
from urllib import urlopen
45
- string_types = basestring
46
+ string_types = basestring # noqa
46
47
else :
47
48
from urllib .request import urlopen
48
49
string_types = str
@@ -70,7 +71,7 @@ def _download_file(baseline, filename):
70
71
71
72
72
73
def pytest_addoption (parser ):
73
- group = parser .getgroup ("general " )
74
+ group = parser .getgroup ("matplotlib image comparison " )
74
75
group .addoption ('--mpl' , action = 'store_true' ,
75
76
help = "Enable comparison of matplotlib figures to reference files" )
76
77
group .addoption ('--mpl-generate-path' ,
@@ -117,6 +118,10 @@ def pytest_configure(config):
117
118
generate_dir = generate_dir ,
118
119
results_dir = results_dir ))
119
120
121
+ else :
122
+
123
+ config .pluginmanager .register (FigureCloser (config ))
124
+
120
125
121
126
@contextlib .contextmanager
122
127
def switch_backend (backend ):
@@ -143,8 +148,14 @@ def __init__(self, config, baseline_dir=None, generate_dir=None, results_dir=Non
143
148
144
149
def pytest_runtest_setup (self , item ):
145
150
151
+ compare = item .keywords .get ('mpl_image_compare' )
152
+
153
+ if compare is None :
154
+ return
155
+
146
156
import matplotlib
147
157
import matplotlib .pyplot as plt
158
+ from matplotlib .figure import Figure
148
159
from matplotlib .testing .compare import compare_images
149
160
from matplotlib .testing .decorators import ImageComparisonTest as MplImageComparisonTest
150
161
try :
@@ -154,11 +165,6 @@ def pytest_runtest_setup(self, item):
154
165
155
166
MPL_LT_15 = LooseVersion (matplotlib .__version__ ) < LooseVersion ('1.5' )
156
167
157
- compare = item .keywords .get ('mpl_image_compare' )
158
-
159
- if compare is None :
160
- return
161
-
162
168
tolerance = compare .kwargs .get ('tolerance' , 2 )
163
169
savefig_kwargs = compare .kwargs .get ('savefig_kwargs' , {})
164
170
style = compare .kwargs .get ('style' , 'classic' )
@@ -188,7 +194,6 @@ def item_function_wrapper(*args, **kwargs):
188
194
with plt .style .context (style ), switch_backend (backend ):
189
195
190
196
# Run test and get figure object
191
- import inspect
192
197
if inspect .ismethod (original ): # method
193
198
# In some cases, for example if setup_method is used,
194
199
# original appears to belong to an instance of the test
@@ -219,7 +224,13 @@ def item_function_wrapper(*args, **kwargs):
219
224
test_image = os .path .abspath (os .path .join (result_dir , filename ))
220
225
221
226
fig .savefig (test_image , ** savefig_kwargs )
222
- plt .close (fig )
227
+
228
+ # We only need to close actual Matplotlib figure objects. If
229
+ # we are dealing with a figure-like object that provides
230
+ # savefig but is not a real Matplotlib object, we shouldn't
231
+ # try closing it here.
232
+ if isinstance (fig , Figure ):
233
+ plt .close (fig )
223
234
224
235
# Find path to baseline image
225
236
if baseline_remote :
@@ -257,3 +268,45 @@ def item_function_wrapper(*args, **kwargs):
257
268
setattr (item .cls , item .function .__name__ , item_function_wrapper )
258
269
else :
259
270
item .obj = item_function_wrapper
271
+
272
+
273
+ class FigureCloser (object ):
274
+ """
275
+ This is used in place of ImageComparison when the --mpl option is not used,
276
+ to make sure that we still close figures returned by tests.
277
+ """
278
+
279
+ def __init__ (self , config ):
280
+ self .config = config
281
+
282
+ def pytest_runtest_setup (self , item ):
283
+
284
+ compare = item .keywords .get ('mpl_image_compare' )
285
+
286
+ if compare is None :
287
+ return
288
+
289
+ import matplotlib .pyplot as plt
290
+ from matplotlib .figure import Figure
291
+
292
+ original = item .function
293
+
294
+ @wraps (item .function )
295
+ def item_function_wrapper (* args , ** kwargs ):
296
+
297
+ if inspect .ismethod (original ): # method
298
+ fig = original .__func__ (* args , ** kwargs )
299
+ else : # function
300
+ fig = original (* args , ** kwargs )
301
+
302
+ # We only need to close actual Matplotlib figure objects. If
303
+ # we are dealing with a figure-like object that provides
304
+ # savefig but is not a real Matplotlib object, we shouldn't
305
+ # try closing it here.
306
+ if isinstance (fig , Figure ):
307
+ plt .close (fig )
308
+
309
+ if item .cls is not None :
310
+ setattr (item .cls , item .function .__name__ , item_function_wrapper )
311
+ else :
312
+ item .obj = item_function_wrapper
0 commit comments