TST: use with where possible instead of manual close

Coincidentally fixes some StataReaders being left open in tests.

Author: akx

pandas/tests/io/parser/test_encoding.py

@@ -66,13 +66,9 @@ def test_utf16_bom_skiprows(all_parsers, sep, encoding):
         with open(path, "wb") as f:
             f.write(bytes_data)
 
-        bytes_buffer = BytesIO(data.encode(utf8))
-        bytes_buffer = TextIOWrapper(bytes_buffer, encoding=utf8)
-
-        result = parser.read_csv(path, encoding=encoding, **kwargs)
-        expected = parser.read_csv(bytes_buffer, encoding=utf8, **kwargs)
-
-        bytes_buffer.close()
+        with TextIOWrapper(BytesIO(data.encode(utf8)), encoding=utf8) as bytes_buffer:
+            result = parser.read_csv(path, encoding=encoding, **kwargs)
+            expected = parser.read_csv(bytes_buffer, encoding=utf8, **kwargs)
         tm.assert_frame_equal(result, expected)
 
 

pandas/tests/io/pytables/test_file_handling.py

@@ -42,9 +42,8 @@ def test_mode(setup_path, mode):
                 HDFStore(path, mode=mode)
 
         else:
-            store = HDFStore(path, mode=mode)
-            assert store._handle.mode == mode
-            store.close()
+            with HDFStore(path, mode=mode) as store:
+                assert store._handle.mode == mode
 
     with ensure_clean_path(setup_path) as path:
 
@@ -253,14 +252,13 @@ def test_complibs(setup_path):
 
             # Open file and check metadata
             # for correct amount of compression
-            h5table = tables.open_file(tmpfile, mode="r")
-            for node in h5table.walk_nodes(where="/" + gname, classname="Leaf"):
-                assert node.filters.complevel == lvl
-                if lvl == 0:
-                    assert node.filters.complib is None
-                else:
-                    assert node.filters.complib == lib
-            h5table.close()
+            with tables.open_file(tmpfile, mode="r") as h5table:
+                for node in h5table.walk_nodes(where="/" + gname, classname="Leaf"):
+                    assert node.filters.complevel == lvl
+                    if lvl == 0:
+                        assert node.filters.complib is None
+                    else:
+                        assert node.filters.complib == lib
 
 
 @pytest.mark.skipif(
@@ -334,15 +332,13 @@ def test_multiple_open_close(setup_path):
 
         if pytables._table_file_open_policy_is_strict:
             # multiples
-            store1 = HDFStore(path)
             msg = (
                 r"The file [\S]* is already opened\.  Please close it before "
                 r"reopening in write mode\."
             )
-            with pytest.raises(ValueError, match=msg):
-                HDFStore(path)
-
-            store1.close()
+            with HDFStore(path) as store1:
+                with pytest.raises(ValueError, match=msg):
+                    HDFStore(path)
         else:
 
             # multiples

pandas/tests/io/pytables/test_read.py

@@ -1,3 +1,4 @@
+from contextlib import closing
 from pathlib import Path
 import re
 
@@ -208,11 +209,10 @@ def test_read_hdf_open_store(setup_path):
     with ensure_clean_path(setup_path) as path:
         df.to_hdf(path, "df", mode="w")
         direct = read_hdf(path, "df")
-        store = HDFStore(path, mode="r")
-        indirect = read_hdf(store, "df")
-        tm.assert_frame_equal(direct, indirect)
-        assert store.is_open
-        store.close()
+        with HDFStore(path, mode="r") as store:
+            indirect = read_hdf(store, "df")
+            tm.assert_frame_equal(direct, indirect)
+            assert store.is_open
 
 
 def test_read_hdf_iterator(setup_path):
@@ -224,10 +224,10 @@ def test_read_hdf_iterator(setup_path):
         df.to_hdf(path, "df", mode="w", format="t")
         direct = read_hdf(path, "df")
         iterator = read_hdf(path, "df", iterator=True)
-        assert isinstance(iterator, TableIterator)
-        indirect = next(iterator.__iter__())
-        tm.assert_frame_equal(direct, indirect)
-        iterator.store.close()
+        with closing(iterator.store):
+            assert isinstance(iterator, TableIterator)
+            indirect = next(iterator.__iter__())
+            tm.assert_frame_equal(direct, indirect)
 
 
 def test_read_nokey(setup_path):

pandas/tests/io/pytables/test_select.py

@@ -684,10 +684,9 @@ def test_frame_select_complex2():
 
         # scope with list like
         l0 = selection.index.tolist()  # noqa:F841
-        store = HDFStore(hh)
-        result = store.select("df", where="l1=l0")
-        tm.assert_frame_equal(result, expected)
-        store.close()
+        with HDFStore(hh) as store:
+            result = store.select("df", where="l1=l0")
+            tm.assert_frame_equal(result, expected)
 
         result = read_hdf(hh, "df", where="l1=l0")
         tm.assert_frame_equal(result, expected)
@@ -707,21 +706,19 @@ def test_frame_select_complex2():
         tm.assert_frame_equal(result, expected)
 
         # scope with index
-        store = HDFStore(hh)
-
-        result = store.select("df", where="l1=index")
-        tm.assert_frame_equal(result, expected)
+        with HDFStore(hh) as store:
 
-        result = store.select("df", where="l1=selection.index")
-        tm.assert_frame_equal(result, expected)
+            result = store.select("df", where="l1=index")
+            tm.assert_frame_equal(result, expected)
 
-        result = store.select("df", where="l1=selection.index.tolist()")
-        tm.assert_frame_equal(result, expected)
+            result = store.select("df", where="l1=selection.index")
+            tm.assert_frame_equal(result, expected)
 
-        result = store.select("df", where="l1=list(selection.index)")
-        tm.assert_frame_equal(result, expected)
+            result = store.select("df", where="l1=selection.index.tolist()")
+            tm.assert_frame_equal(result, expected)
 
-        store.close()
+            result = store.select("df", where="l1=list(selection.index)")
+            tm.assert_frame_equal(result, expected)
 
 
 def test_invalid_filtering(setup_path):

pandas/tests/io/pytables/test_store.py

@@ -918,9 +918,8 @@ def do_copy(f, new_f=None, keys=None, propindexes=True, **kwargs):
         df = tm.makeDataFrame()
 
         with tm.ensure_clean() as path:
-            st = HDFStore(path)
-            st.append("df", df, data_columns=["A"])
-            st.close()
+            with HDFStore(path) as st:
+                st.append("df", df, data_columns=["A"])
             do_copy(f=path)
             do_copy(f=path, propindexes=False)
 

pandas/tests/io/test_common.py

@@ -118,11 +118,11 @@ def test_get_handle_with_path(self, path_type):
                 assert os.path.expanduser(filename) == handles.handle.name
 
     def test_get_handle_with_buffer(self):
-        input_buffer = StringIO()
-        with icom.get_handle(input_buffer, "r") as handles:
-            assert handles.handle == input_buffer
-        assert not input_buffer.closed
-        input_buffer.close()
+        with StringIO() as input_buffer:
+            with icom.get_handle(input_buffer, "r") as handles:
+                assert handles.handle == input_buffer
+            assert not input_buffer.closed
+        assert input_buffer.closed
 
     # Test that BytesIOWrapper(get_handle) returns correct amount of bytes every time
     def test_bytesiowrapper_returns_correct_bytes(self):

pandas/tests/io/test_compression.py

@@ -282,18 +282,17 @@ def test_bzip_compression_level(obj, method):
 )
 def test_empty_archive_zip(suffix, archive):
     with tm.ensure_clean(filename=suffix) as path:
-        file = archive(path, "w")
-        file.close()
+        with archive(path, "w"):
+            pass
         with pytest.raises(ValueError, match="Zero files found"):
             pd.read_csv(path)
 
 
 def test_ambiguous_archive_zip():
     with tm.ensure_clean(filename=".zip") as path:
-        file = zipfile.ZipFile(path, "w")
-        file.writestr("a.csv", "foo,bar")
-        file.writestr("b.csv", "foo,bar")
-        file.close()
+        with zipfile.ZipFile(path, "w") as file:
+            file.writestr("a.csv", "foo,bar")
+            file.writestr("b.csv", "foo,bar")
         with pytest.raises(ValueError, match="Multiple files found in ZIP file"):
             pd.read_csv(path)
 

pandas/tests/io/test_fsspec.py

@@ -95,21 +95,19 @@ def test_to_csv_fsspec_object(cleared_fs, binary_mode, df1):
 
     path = "memory://test/test.csv"
     mode = "wb" if binary_mode else "w"
-    fsspec_object = fsspec.open(path, mode=mode).open()
-
-    df1.to_csv(fsspec_object, index=True)
-    assert not fsspec_object.closed
+    with fsspec.open(path, mode=mode).open() as fsspec_object:
+        df1.to_csv(fsspec_object, index=True)
+        assert not fsspec_object.closed
     fsspec_object.close()
 
     mode = mode.replace("w", "r")
-    fsspec_object = fsspec.open(path, mode=mode).open()
-
-    df2 = read_csv(
-        fsspec_object,
-        parse_dates=["dt"],
-        index_col=0,
-    )
-    assert not fsspec_object.closed
+    with fsspec.open(path, mode=mode) as fsspec_object:
+        df2 = read_csv(
+            fsspec_object,
+            parse_dates=["dt"],
+            index_col=0,
+        )
+        assert not fsspec_object.closed
     fsspec_object.close()
 
     tm.assert_frame_equal(df1, df2)

pandas/tests/io/test_sql.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import csv
+from contextlib import closing
 from datetime import (
     date,
     datetime,
@@ -455,9 +456,8 @@ def sqlite_iris_conn(sqlite_iris_engine):
 
 @pytest.fixture
 def sqlite_buildin():
-    conn = sqlite3.connect(":memory:")
-    yield conn
-    conn.close()
+    with sqlite3.connect(":memory:") as conn:
+        yield conn
 
 
 @pytest.fixture
@@ -1532,13 +1532,14 @@ def test_sql_open_close(self, test_frame3):
 
         with tm.ensure_clean() as name:
 
-            conn = self.connect(name)
-            assert sql.to_sql(test_frame3, "test_frame3_legacy", conn, index=False) == 4
-            conn.close()
+            with closing(self.connect(name)) as conn:
+                assert (
+                    sql.to_sql(test_frame3, "test_frame3_legacy", conn, index=False)
+                    == 4
+                )
 
-            conn = self.connect(name)
-            result = sql.read_sql_query("SELECT * FROM test_frame3_legacy;", conn)
-            conn.close()
+            with closing(self.connect(name)) as conn:
+                result = sql.read_sql_query("SELECT * FROM test_frame3_legacy;", conn)
 
         tm.assert_frame_equal(test_frame3, result)
 
@@ -2371,18 +2372,15 @@ class Test(BaseModel):
 
         BaseModel.metadata.create_all(self.conn)
         Session = sessionmaker(bind=self.conn)
-        session = Session()
-
-        df = DataFrame({"id": [0, 1], "foo": ["hello", "world"]})
-        assert (
-            df.to_sql("test_frame", con=self.conn, index=False, if_exists="replace")
-            == 2
-        )
-
-        session.commit()
-        foo = session.query(Test.id, Test.foo)
-        df = DataFrame(foo)
-        session.close()
+        with Session() as session:
+            df = DataFrame({"id": [0, 1], "foo": ["hello", "world"]})
+            assert (
+                df.to_sql("test_frame", con=self.conn, index=False, if_exists="replace")
+                == 2
+            )
+            session.commit()
+            foo = session.query(Test.id, Test.foo)
+            df = DataFrame(foo)
 
         assert list(df.columns) == ["id", "foo"]