cleanup should not always create a generator.

anntzer · anntzer · commit b441be2a3ea4 · 2016-01-08T08:57:24.000-08:00
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -130,22 +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:
-                if inspect.isgeneratorfunction(func):
+        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
-                else:
+                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)
+                finally:
+                    _do_cleanup(original_units_registry,
+                                original_settings)
 
-        return wrapped_function
+        return wrapped_callable
 
     if isinstance(style, six.string_types):
         return make_cleanup
