Skip to content

Commit b48a73f

Browse files
authored
TST: use with where possible instead of manual close (#48931)
Coincidentally fixes some StataReaders being left open in tests.
1 parent cdb905a commit b48a73f

10 files changed

+127
-151
lines changed

pandas/tests/io/parser/test_encoding.py

+3-7
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,9 @@ def test_utf16_bom_skiprows(all_parsers, sep, encoding):
6666
with open(path, "wb") as f:
6767
f.write(bytes_data)
6868

69-
bytes_buffer = BytesIO(data.encode(utf8))
70-
bytes_buffer = TextIOWrapper(bytes_buffer, encoding=utf8)
71-
72-
result = parser.read_csv(path, encoding=encoding, **kwargs)
73-
expected = parser.read_csv(bytes_buffer, encoding=utf8, **kwargs)
74-
75-
bytes_buffer.close()
69+
with TextIOWrapper(BytesIO(data.encode(utf8)), encoding=utf8) as bytes_buffer:
70+
result = parser.read_csv(path, encoding=encoding, **kwargs)
71+
expected = parser.read_csv(bytes_buffer, encoding=utf8, **kwargs)
7672
tm.assert_frame_equal(result, expected)
7773

7874

pandas/tests/io/pytables/test_file_handling.py

+10-13
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,8 @@ def test_mode(setup_path, tmp_path, mode):
4141
HDFStore(path, mode=mode)
4242

4343
else:
44-
store = HDFStore(path, mode=mode)
45-
assert store._handle.mode == mode
46-
store.close()
44+
with HDFStore(path, mode=mode) as store:
45+
assert store._handle.mode == mode
4746

4847
path = tmp_path / setup_path
4948

@@ -253,16 +252,14 @@ def test_complibs(tmp_path, setup_path):
253252
result = read_hdf(tmpfile, gname)
254253
tm.assert_frame_equal(result, df)
255254

256-
# Open file and check metadata
257-
# for correct amount of compression
258-
h5table = tables.open_file(tmpfile, mode="r")
259-
for node in h5table.walk_nodes(where="/" + gname, classname="Leaf"):
260-
assert node.filters.complevel == lvl
261-
if lvl == 0:
262-
assert node.filters.complib is None
263-
else:
264-
assert node.filters.complib == lib
265-
h5table.close()
255+
# Open file and check metadata for correct amount of compression
256+
with tables.open_file(tmpfile, mode="r") as h5table:
257+
for node in h5table.walk_nodes(where="/" + gname, classname="Leaf"):
258+
assert node.filters.complevel == lvl
259+
if lvl == 0:
260+
assert node.filters.complib is None
261+
else:
262+
assert node.filters.complib == lib
266263

267264

268265
@pytest.mark.skipif(

pandas/tests/io/pytables/test_read.py

+9-9
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from contextlib import closing
12
from pathlib import Path
23
import re
34

@@ -207,11 +208,10 @@ def test_read_hdf_open_store(tmp_path, setup_path):
207208
path = tmp_path / setup_path
208209
df.to_hdf(path, "df", mode="w")
209210
direct = read_hdf(path, "df")
210-
store = HDFStore(path, mode="r")
211-
indirect = read_hdf(store, "df")
212-
tm.assert_frame_equal(direct, indirect)
213-
assert store.is_open
214-
store.close()
211+
with HDFStore(path, mode="r") as store:
212+
indirect = read_hdf(store, "df")
213+
tm.assert_frame_equal(direct, indirect)
214+
assert store.is_open
215215

216216

217217
def test_read_hdf_iterator(tmp_path, setup_path):
@@ -223,10 +223,10 @@ def test_read_hdf_iterator(tmp_path, setup_path):
223223
df.to_hdf(path, "df", mode="w", format="t")
224224
direct = read_hdf(path, "df")
225225
iterator = read_hdf(path, "df", iterator=True)
226-
assert isinstance(iterator, TableIterator)
227-
indirect = next(iterator.__iter__())
228-
tm.assert_frame_equal(direct, indirect)
229-
iterator.store.close()
226+
with closing(iterator.store):
227+
assert isinstance(iterator, TableIterator)
228+
indirect = next(iterator.__iter__())
229+
tm.assert_frame_equal(direct, indirect)
230230

231231

232232
def test_read_nokey(tmp_path, setup_path):

pandas/tests/io/pytables/test_select.py

+12-16
Original file line numberDiff line numberDiff line change
@@ -682,10 +682,9 @@ def test_frame_select_complex2(tmp_path):
682682

683683
# scope with list like
684684
l0 = selection.index.tolist() # noqa:F841
685-
store = HDFStore(hh)
686-
result = store.select("df", where="l1=l0")
687-
tm.assert_frame_equal(result, expected)
688-
store.close()
685+
with HDFStore(hh) as store:
686+
result = store.select("df", where="l1=l0")
687+
tm.assert_frame_equal(result, expected)
689688

690689
result = read_hdf(hh, "df", where="l1=l0")
691690
tm.assert_frame_equal(result, expected)
@@ -705,21 +704,18 @@ def test_frame_select_complex2(tmp_path):
705704
tm.assert_frame_equal(result, expected)
706705

707706
# scope with index
708-
store = HDFStore(hh)
709-
710-
result = store.select("df", where="l1=index")
711-
tm.assert_frame_equal(result, expected)
712-
713-
result = store.select("df", where="l1=selection.index")
714-
tm.assert_frame_equal(result, expected)
707+
with HDFStore(hh) as store:
708+
result = store.select("df", where="l1=index")
709+
tm.assert_frame_equal(result, expected)
715710

716-
result = store.select("df", where="l1=selection.index.tolist()")
717-
tm.assert_frame_equal(result, expected)
711+
result = store.select("df", where="l1=selection.index")
712+
tm.assert_frame_equal(result, expected)
718713

719-
result = store.select("df", where="l1=list(selection.index)")
720-
tm.assert_frame_equal(result, expected)
714+
result = store.select("df", where="l1=selection.index.tolist()")
715+
tm.assert_frame_equal(result, expected)
721716

722-
store.close()
717+
result = store.select("df", where="l1=list(selection.index)")
718+
tm.assert_frame_equal(result, expected)
723719

724720

725721
def test_invalid_filtering(setup_path):

pandas/tests/io/pytables/test_store.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -917,9 +917,8 @@ def do_copy(f, new_f=None, keys=None, propindexes=True, **kwargs):
917917
df = tm.makeDataFrame()
918918

919919
with tm.ensure_clean() as path:
920-
st = HDFStore(path)
921-
st.append("df", df, data_columns=["A"])
922-
st.close()
920+
with HDFStore(path) as st:
921+
st.append("df", df, data_columns=["A"])
923922
do_copy(f=path)
924923
do_copy(f=path, propindexes=False)
925924

pandas/tests/io/test_common.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,11 @@ def test_get_handle_with_path(self, path_type):
118118
assert os.path.expanduser(filename) == handles.handle.name
119119

120120
def test_get_handle_with_buffer(self):
121-
input_buffer = StringIO()
122-
with icom.get_handle(input_buffer, "r") as handles:
123-
assert handles.handle == input_buffer
124-
assert not input_buffer.closed
125-
input_buffer.close()
121+
with StringIO() as input_buffer:
122+
with icom.get_handle(input_buffer, "r") as handles:
123+
assert handles.handle == input_buffer
124+
assert not input_buffer.closed
125+
assert input_buffer.closed
126126

127127
# Test that BytesIOWrapper(get_handle) returns correct amount of bytes every time
128128
def test_bytesiowrapper_returns_correct_bytes(self):

pandas/tests/io/test_compression.py

+5-6
Original file line numberDiff line numberDiff line change
@@ -282,18 +282,17 @@ def test_bzip_compression_level(obj, method):
282282
)
283283
def test_empty_archive_zip(suffix, archive):
284284
with tm.ensure_clean(filename=suffix) as path:
285-
file = archive(path, "w")
286-
file.close()
285+
with archive(path, "w"):
286+
pass
287287
with pytest.raises(ValueError, match="Zero files found"):
288288
pd.read_csv(path)
289289

290290

291291
def test_ambiguous_archive_zip():
292292
with tm.ensure_clean(filename=".zip") as path:
293-
file = zipfile.ZipFile(path, "w")
294-
file.writestr("a.csv", "foo,bar")
295-
file.writestr("b.csv", "foo,bar")
296-
file.close()
293+
with zipfile.ZipFile(path, "w") as file:
294+
file.writestr("a.csv", "foo,bar")
295+
file.writestr("b.csv", "foo,bar")
297296
with pytest.raises(ValueError, match="Multiple files found in ZIP file"):
298297
pd.read_csv(path)
299298

pandas/tests/io/test_fsspec.py

+10-14
Original file line numberDiff line numberDiff line change
@@ -95,22 +95,18 @@ def test_to_csv_fsspec_object(cleared_fs, binary_mode, df1):
9595

9696
path = "memory://test/test.csv"
9797
mode = "wb" if binary_mode else "w"
98-
fsspec_object = fsspec.open(path, mode=mode).open()
99-
100-
df1.to_csv(fsspec_object, index=True)
101-
assert not fsspec_object.closed
102-
fsspec_object.close()
98+
with fsspec.open(path, mode=mode).open() as fsspec_object:
99+
df1.to_csv(fsspec_object, index=True)
100+
assert not fsspec_object.closed
103101

104102
mode = mode.replace("w", "r")
105-
fsspec_object = fsspec.open(path, mode=mode).open()
106-
107-
df2 = read_csv(
108-
fsspec_object,
109-
parse_dates=["dt"],
110-
index_col=0,
111-
)
112-
assert not fsspec_object.closed
113-
fsspec_object.close()
103+
with fsspec.open(path, mode=mode) as fsspec_object:
104+
df2 = read_csv(
105+
fsspec_object,
106+
parse_dates=["dt"],
107+
index_col=0,
108+
)
109+
assert not fsspec_object.closed
114110

115111
tm.assert_frame_equal(df1, df2)
116112

pandas/tests/io/test_sql.py

+19-21
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
"""
1919
from __future__ import annotations
2020

21+
from contextlib import closing
2122
import csv
2223
from datetime import (
2324
date,
@@ -455,9 +456,8 @@ def sqlite_iris_conn(sqlite_iris_engine):
455456

456457
@pytest.fixture
457458
def sqlite_buildin():
458-
conn = sqlite3.connect(":memory:")
459-
yield conn
460-
conn.close()
459+
with sqlite3.connect(":memory:") as conn:
460+
yield conn
461461

462462

463463
@pytest.fixture
@@ -1532,13 +1532,14 @@ def test_sql_open_close(self, test_frame3):
15321532

15331533
with tm.ensure_clean() as name:
15341534

1535-
conn = self.connect(name)
1536-
assert sql.to_sql(test_frame3, "test_frame3_legacy", conn, index=False) == 4
1537-
conn.close()
1535+
with closing(self.connect(name)) as conn:
1536+
assert (
1537+
sql.to_sql(test_frame3, "test_frame3_legacy", conn, index=False)
1538+
== 4
1539+
)
15381540

1539-
conn = self.connect(name)
1540-
result = sql.read_sql_query("SELECT * FROM test_frame3_legacy;", conn)
1541-
conn.close()
1541+
with closing(self.connect(name)) as conn:
1542+
result = sql.read_sql_query("SELECT * FROM test_frame3_legacy;", conn)
15421543

15431544
tm.assert_frame_equal(test_frame3, result)
15441545

@@ -2371,18 +2372,15 @@ class Test(BaseModel):
23712372

23722373
BaseModel.metadata.create_all(self.conn)
23732374
Session = sessionmaker(bind=self.conn)
2374-
session = Session()
2375-
2376-
df = DataFrame({"id": [0, 1], "foo": ["hello", "world"]})
2377-
assert (
2378-
df.to_sql("test_frame", con=self.conn, index=False, if_exists="replace")
2379-
== 2
2380-
)
2381-
2382-
session.commit()
2383-
foo = session.query(Test.id, Test.foo)
2384-
df = DataFrame(foo)
2385-
session.close()
2375+
with Session() as session:
2376+
df = DataFrame({"id": [0, 1], "foo": ["hello", "world"]})
2377+
assert (
2378+
df.to_sql("test_frame", con=self.conn, index=False, if_exists="replace")
2379+
== 2
2380+
)
2381+
session.commit()
2382+
foo = session.query(Test.id, Test.foo)
2383+
df = DataFrame(foo)
23862384

23872385
assert list(df.columns) == ["id", "foo"]
23882386

0 commit comments

Comments
 (0)