use ensure_clean rather than explicit os.remove #34384 (#34385)

Author: jnecus

pandas/tests/io/excel/test_openpyxl.py

@@ -1,5 +1,3 @@
-import os
-
 import numpy as np
 import pytest
 
@@ -107,17 +105,15 @@ def test_write_append_mode(ext, mode, expected):
             assert wb2.worksheets[index]["A1"].value == cell_value
 
 
-def test_to_excel_with_openpyxl_engine(ext, tmpdir):
+def test_to_excel_with_openpyxl_engine(ext):
     # GH 29854
-    df1 = DataFrame({"A": np.linspace(1, 10, 10)})
-    df2 = DataFrame({"B": np.linspace(1, 20, 10)})
-    df = pd.concat([df1, df2], axis=1)
-    styled = df.style.applymap(
-        lambda val: "color: %s" % ("red" if val < 0 else "black")
-    ).highlight_max()
-
-    filename = tmpdir / "styled.xlsx"
-    styled.to_excel(filename, engine="openpyxl")
-
-    assert filename.exists()
-    os.remove(filename)
+    with tm.ensure_clean(ext) as filename:
+
+        df1 = DataFrame({"A": np.linspace(1, 10, 10)})
+        df2 = DataFrame({"B": np.linspace(1, 20, 10)})
+        df = pd.concat([df1, df2], axis=1)
+        styled = df.style.applymap(
+            lambda val: "color: %s" % ("red" if val < 0 else "black")
+        ).highlight_max()
+
+        styled.to_excel(filename, engine="openpyxl")

pandas/tests/io/pytables/test_store.py

@@ -36,11 +36,9 @@
 import pandas._testing as tm
 from pandas.tests.io.pytables.common import (
     _maybe_remove,
-    create_tempfile,
     ensure_clean_path,
     ensure_clean_store,
     safe_close,
-    safe_remove,
     tables,
 )
 
@@ -80,33 +78,25 @@ def test_format_kwarg_in_constructor(self, setup_path):
 
         msg = "format is not a defined argument for HDFStore"
 
-        with ensure_clean_path(setup_path) as path:
+        with tm.ensure_clean(setup_path) as path:
             with pytest.raises(ValueError, match=msg):
                 HDFStore(path, format="table")
 
     def test_context(self, setup_path):
-        path = create_tempfile(setup_path)
-        try:
-            with HDFStore(path) as tbl:
-                raise ValueError("blah")
-        except ValueError:
-            pass
-        finally:
-            safe_remove(path)
-
-        try:
+        with tm.ensure_clean(setup_path) as path:
+            try:
+                with HDFStore(path) as tbl:
+                    raise ValueError("blah")
+            except ValueError:
+                pass
+        with tm.ensure_clean(setup_path) as path:
             with HDFStore(path) as tbl:
                 tbl["a"] = tm.makeDataFrame()
-
-            with HDFStore(path) as tbl:
                 assert len(tbl) == 1
                 assert type(tbl["a"]) == DataFrame
-        finally:
-            safe_remove(path)
 
     def test_conv_read_write(self, setup_path):
-        path = create_tempfile(setup_path)
-        try:
+        with tm.ensure_clean() as path:
 
             def roundtrip(key, obj, **kwargs):
                 obj.to_hdf(path, key, **kwargs)
@@ -127,9 +117,6 @@ def roundtrip(key, obj, **kwargs):
             result = read_hdf(path, "table", where=["index>2"])
             tm.assert_frame_equal(df[df.index > 2], result)
 
-        finally:
-            safe_remove(path)
-
     def test_long_strings(self, setup_path):
 
         # GH6166
@@ -605,7 +592,7 @@ def test_reopen_handle(self, setup_path):
 
     def test_open_args(self, setup_path):
 
-        with ensure_clean_path(setup_path) as path:
+        with tm.ensure_clean(setup_path) as path:
 
             df = tm.makeDataFrame()
 
@@ -621,8 +608,8 @@ def test_open_args(self, setup_path):
 
             store.close()
 
-            # the file should not have actually been written
-            assert not os.path.exists(path)
+        # the file should not have actually been written
+        assert not os.path.exists(path)
 
     def test_flush(self, setup_path):
 
@@ -4194,7 +4181,6 @@ def do_copy(f, new_f=None, keys=None, propindexes=True, **kwargs):
                         import tempfile
 
                         fd, new_f = tempfile.mkstemp()
-
                     tstore = store.copy(
                         new_f, keys=keys, propindexes=propindexes, **kwargs
                     )
@@ -4225,20 +4211,17 @@ def do_copy(f, new_f=None, keys=None, propindexes=True, **kwargs):
                         os.close(fd)
                     except (OSError, ValueError):
                         pass
-                    safe_remove(new_f)
+                    os.remove(new_f)
 
             # new table
             df = tm.makeDataFrame()
 
-            try:
-                path = create_tempfile(setup_path)
+            with tm.ensure_clean() as path:
                 st = HDFStore(path)
                 st.append("df", df, data_columns=["A"])
                 st.close()
                 do_copy(f=path)
                 do_copy(f=path, propindexes=False)
-            finally:
-                safe_remove(path)
 
     def test_store_datetime_fractional_secs(self, setup_path):
 
@@ -4677,7 +4660,7 @@ def test_read_hdf_generic_buffer_errors(self):
 
     def test_invalid_complib(self, setup_path):
         df = DataFrame(np.random.rand(4, 5), index=list("abcd"), columns=list("ABCDE"))
-        with ensure_clean_path(setup_path) as path:
+        with tm.ensure_clean(setup_path) as path:
             with pytest.raises(ValueError):
                 df.to_hdf(path, "df", complib="foolib")