Merge pull request #5809 from anntzer/cleanup-generative-tests

mdboom · tacaswell · commit 0375ea4ff194 · 2016-03-26T18:33:49.000-04:00
Support generative tests in @cleanup.
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -5,6 +5,7 @@
 
 import functools
 import gc
+import inspect
 import os
 import sys
 import shutil
@@ -129,18 +130,31 @@ def cleanup(style=None):
     # writing a decorator with optional arguments.
 
     def make_cleanup(func):
-        @functools.wraps(func)
-        def wrapped_function(*args, **kwargs):
-            original_units_registry = matplotlib.units.registry.copy()
-            original_settings = mpl.rcParams.copy()
-            matplotlib.style.use(style)
-            try:
-                func(*args, **kwargs)
-            finally:
-                _do_cleanup(original_units_registry,
-                            original_settings)
+        if inspect.isgenerator(func):
+            @functools.wraps(func)
+            def wrapped_callable(*args, **kwargs):
+                original_units_registry = matplotlib.units.registry.copy()
+                original_settings = mpl.rcParams.copy()
+                matplotlib.style.use(style)
+                try:
+                    for yielded in func(*args, **kwargs):
+                        yield yielded
+                finally:
+                    _do_cleanup(original_units_registry,
+                                original_settings)
+        else:
+            @functools.wraps(func)
+            def wrapped_callable(*args, **kwargs):
+                original_units_registry = matplotlib.units.registry.copy()
+                original_settings = mpl.rcParams.copy()
+                matplotlib.style.use(style)
+                try:
+                    func(*args, **kwargs)
+                finally:
+                    _do_cleanup(original_units_registry,
+                                original_settings)
 
-        return wrapped_function
+        return wrapped_callable
 
     if isinstance(style, six.string_types):
         return make_cleanup
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -4282,7 +4282,7 @@ def _helper_y(ax):
     orig_xlim = ax_lst[0][1].get_xlim()
     ax.remove()
     ax.set_xlim(0, 5)
-    assert assert_array_equal(ax_lst[0][1].get_xlim(), orig_xlim)
+    assert_array_equal(ax_lst[0][1].get_xlim(), orig_xlim)
 
 
 @cleanup
