Merge pull request #5359 from jtratner/always-cleanup-graphics-files

TST: Make sure files are removed in test_graphics

Author: jtratner

doc/source/release.rst

@@ -753,6 +753,8 @@ Bug Fixes
   - Fix return value/type signature of ``initObjToJSON()`` to be compatible
     with numpy's ``import_array()`` (:issue:`5334`, :issue:`5326`)
   - Bug when renaming then set_index on a DataFrame (:issue:`5344`)
+  - Test suite no longer leaves around temporary files when testing graphics. (:issue:`5347`)
+    (thanks for catching this @yarikoptic!)
 
 pandas 0.12.0
 -------------

pandas/tests/test_graphics.py

@@ -1163,7 +1163,7 @@ def _check_plot_works(f, *args, **kwargs):
         else:
             assert_is_valid_plot_return_object(ret)
 
-        with ensure_clean() as path:
+        with ensure_clean(return_filelike=True) as path:
             plt.savefig(path)
     finally:
         tm.close(fig)

pandas/tseries/tests/test_plotting.py

@@ -943,7 +943,7 @@ def _check_plot_works(f, freq=None, series=None, *args, **kwargs):
         except Exception:
             pass
 
-        with ensure_clean() as path:
+        with ensure_clean(return_filelike=True) as path:
             plt.savefig(path)
     finally:
         plt.close(fig)

pandas/util/testing.py

@@ -306,18 +306,39 @@ def set_trace():
 
 
 @contextmanager
-def ensure_clean(filename=None):
-    # if we are not passed a filename, generate a temporary
-    if filename is None:
-        filename = tempfile.mkstemp()[1]
+def ensure_clean(filename=None, return_filelike=False):
+    """Gets a temporary path and agrees to remove on close.
 
-    try:
-        yield filename
-    finally:
+    Parameters
+    ----------
+    filename : str (optional)
+        if None, creates a temporary file which is then removed when out of
+        scope.
+    return_filelike: bool (default False)
+        if True, returns a file-like which is *always* cleaned. Necessary for
+        savefig and other functions which want to append extensions. Ignores
+        filename if True.
+    """
+
+    if return_filelike:
+        f = tempfile.TemporaryFile()
         try:
-            os.remove(filename)
-        except:
-            pass
+            yield f
+        finally:
+            f.close()
+
+    else:
+        # if we are not passed a filename, generate a temporary
+        if filename is None:
+            filename = tempfile.mkstemp()[1]
+
+        try:
+            yield filename
+        finally:
+            try:
+                os.remove(filename)
+            except Exception as e:
+                print(e)
 
 
 def get_data_path(f=''):