44import six
55
66import functools
7+ import gc
78import os
89import sys
910import shutil
@@ -65,20 +66,25 @@ def failer(*args, **kwargs):
6566 return known_fail_decorator
6667
6768
69+ def _do_cleanup (original_units_registry ):
70+ plt .close ('all' )
71+ gc .collect ()
72+
73+ matplotlib .tests .setup ()
74+
75+ matplotlib .units .registry .clear ()
76+ matplotlib .units .registry .update (original_units_registry )
77+ warnings .resetwarnings () # reset any warning filters set in tests
78+
79+
6880class CleanupTest (object ):
6981 @classmethod
7082 def setup_class (cls ):
7183 cls .original_units_registry = matplotlib .units .registry .copy ()
7284
7385 @classmethod
7486 def teardown_class (cls ):
75- plt .close ('all' )
76-
77- matplotlib .tests .setup ()
78-
79- matplotlib .units .registry .clear ()
80- matplotlib .units .registry .update (cls .original_units_registry )
81- warnings .resetwarnings () # reset any warning filters set in tests
87+ _do_cleanup (cls .original_units_registry )
8288
8389 def test (self ):
8490 self ._func ()
@@ -93,13 +99,7 @@ def setUpClass(cls):
9399
94100 @classmethod
95101 def tearDownClass (cls ):
96- plt .close ('all' )
97-
98- matplotlib .tests .setup ()
99-
100- matplotlib .units .registry .clear ()
101- matplotlib .units .registry .update (cls .original_units_registry )
102- warnings .resetwarnings () # reset any warning filters set in tests
102+ _do_cleanup (cls .original_units_registry )
103103
104104
105105def cleanup (func ):
@@ -109,13 +109,8 @@ def wrapped_function(*args, **kwargs):
109109 try :
110110 func (* args , ** kwargs )
111111 finally :
112- plt .close ('all' )
113-
114- matplotlib .tests .setup ()
112+ _do_cleanup (original_units_registry )
115113
116- matplotlib .units .registry .clear ()
117- matplotlib .units .registry .update (original_units_registry )
118- warnings .resetwarnings () #reset any warning filters set in tests
119114 return wrapped_function
120115
121116
@@ -157,8 +152,6 @@ def test(self):
157152 baseline_dir , result_dir = _image_directories (self ._func )
158153
159154 for fignum , baseline in zip (plt .get_fignums (), self ._baseline_images ):
160- figure = plt .figure (fignum )
161-
162155 for extension in self ._extensions :
163156 will_fail = not extension in comparable_formats ()
164157 if will_fail :
@@ -182,6 +175,8 @@ def test(self):
182175 will_fail , fail_msg ,
183176 known_exception_class = ImageComparisonFailure )
184177 def do_test ():
178+ figure = plt .figure (fignum )
179+
185180 if self ._remove_text :
186181 self .remove_text (figure )
187182
0 commit comments