Skip to content

Commit fc9fdba

Browse files
authored
TST: use more context managers (#39785)
1 parent bcf2406 commit fc9fdba

File tree

6 files changed

+116
-106
lines changed

6 files changed

+116
-106
lines changed

pandas/tests/io/excel/test_writers.py

Lines changed: 67 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -342,8 +342,8 @@ def test_excel_sheet_by_name_raise(self, path, engine):
342342
gt = DataFrame(np.random.randn(10, 2))
343343
gt.to_excel(path)
344344

345-
xl = ExcelFile(path)
346-
df = pd.read_excel(xl, sheet_name=0, index_col=0)
345+
with ExcelFile(path) as xl:
346+
df = pd.read_excel(xl, sheet_name=0, index_col=0)
347347

348348
tm.assert_frame_equal(gt, df)
349349

@@ -419,8 +419,8 @@ def test_mixed(self, frame, path):
419419
mixed_frame["foo"] = "bar"
420420

421421
mixed_frame.to_excel(path, "test1")
422-
reader = ExcelFile(path)
423-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
422+
with ExcelFile(path) as reader:
423+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
424424
tm.assert_frame_equal(mixed_frame, recons)
425425

426426
def test_ts_frame(self, tsframe, path):
@@ -431,9 +431,8 @@ def test_ts_frame(self, tsframe, path):
431431
df.index = index
432432

433433
df.to_excel(path, "test1")
434-
reader = ExcelFile(path)
435-
436-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
434+
with ExcelFile(path) as reader:
435+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
437436
tm.assert_frame_equal(df, recons)
438437

439438
def test_basics_with_nan(self, frame, path):
@@ -451,8 +450,8 @@ def test_int_types(self, np_type, path):
451450
df = DataFrame(np.random.randint(-10, 10, size=(10, 2)), dtype=np_type)
452451
df.to_excel(path, "test1")
453452

454-
reader = ExcelFile(path)
455-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
453+
with ExcelFile(path) as reader:
454+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
456455

457456
int_frame = df.astype(np.int64)
458457
tm.assert_frame_equal(int_frame, recons)
@@ -475,8 +474,10 @@ def test_float_types(self, np_type, path):
475474
df = DataFrame(np.random.random_sample(10), dtype=np_type)
476475
df.to_excel(path, "test1")
477476

478-
reader = ExcelFile(path)
479-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(np_type)
477+
with ExcelFile(path) as reader:
478+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(
479+
np_type
480+
)
480481

481482
tm.assert_frame_equal(df, recons)
482483

@@ -486,17 +487,19 @@ def test_bool_types(self, np_type, path):
486487
df = DataFrame([1, 0, True, False], dtype=np_type)
487488
df.to_excel(path, "test1")
488489

489-
reader = ExcelFile(path)
490-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(np_type)
490+
with ExcelFile(path) as reader:
491+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(
492+
np_type
493+
)
491494

492495
tm.assert_frame_equal(df, recons)
493496

494497
def test_inf_roundtrip(self, path):
495498
df = DataFrame([(1, np.inf), (2, 3), (5, -np.inf)])
496499
df.to_excel(path, "test1")
497500

498-
reader = ExcelFile(path)
499-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
501+
with ExcelFile(path) as reader:
502+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
500503

501504
tm.assert_frame_equal(df, recons)
502505

@@ -518,11 +521,11 @@ def test_sheets(self, frame, tsframe, path):
518521
with ExcelWriter(path) as writer:
519522
frame.to_excel(writer, "test1")
520523
tsframe.to_excel(writer, "test2")
521-
reader = ExcelFile(path)
522-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
523-
tm.assert_frame_equal(frame, recons)
524-
recons = pd.read_excel(reader, sheet_name="test2", index_col=0)
525-
tm.assert_frame_equal(tsframe, recons)
524+
with ExcelFile(path) as reader:
525+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
526+
tm.assert_frame_equal(frame, recons)
527+
recons = pd.read_excel(reader, sheet_name="test2", index_col=0)
528+
tm.assert_frame_equal(tsframe, recons)
526529
assert 2 == len(reader.sheet_names)
527530
assert "test1" == reader.sheet_names[0]
528531
assert "test2" == reader.sheet_names[1]
@@ -539,8 +542,8 @@ def test_colaliases(self, frame, path):
539542
# column aliases
540543
col_aliases = Index(["AA", "X", "Y", "Z"])
541544
frame.to_excel(path, "test1", header=col_aliases)
542-
reader = ExcelFile(path)
543-
rs = pd.read_excel(reader, sheet_name="test1", index_col=0)
545+
with ExcelFile(path) as reader:
546+
rs = pd.read_excel(reader, sheet_name="test1", index_col=0)
544547
xp = frame.copy()
545548
xp.columns = col_aliases
546549
tm.assert_frame_equal(xp, rs)
@@ -557,8 +560,10 @@ def test_roundtrip_indexlabels(self, merge_cells, frame, path):
557560
# test index_label
558561
df = DataFrame(np.random.randn(10, 2)) >= 0
559562
df.to_excel(path, "test1", index_label=["test"], merge_cells=merge_cells)
560-
reader = ExcelFile(path)
561-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(np.int64)
563+
with ExcelFile(path) as reader:
564+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(
565+
np.int64
566+
)
562567
df.index.names = ["test"]
563568
assert df.index.names == recons.index.names
564569

@@ -569,15 +574,19 @@ def test_roundtrip_indexlabels(self, merge_cells, frame, path):
569574
index_label=["test", "dummy", "dummy2"],
570575
merge_cells=merge_cells,
571576
)
572-
reader = ExcelFile(path)
573-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(np.int64)
577+
with ExcelFile(path) as reader:
578+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(
579+
np.int64
580+
)
574581
df.index.names = ["test"]
575582
assert df.index.names == recons.index.names
576583

577584
df = DataFrame(np.random.randn(10, 2)) >= 0
578585
df.to_excel(path, "test1", index_label="test", merge_cells=merge_cells)
579-
reader = ExcelFile(path)
580-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(np.int64)
586+
with ExcelFile(path) as reader:
587+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0).astype(
588+
np.int64
589+
)
581590
df.index.names = ["test"]
582591
tm.assert_frame_equal(df, recons.astype(bool))
583592

@@ -592,8 +601,8 @@ def test_roundtrip_indexlabels(self, merge_cells, frame, path):
592601
df = frame.copy()
593602
df = df.set_index(["A", "B"])
594603

595-
reader = ExcelFile(path)
596-
recons = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
604+
with ExcelFile(path) as reader:
605+
recons = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
597606
tm.assert_frame_equal(df, recons)
598607

599608
def test_excel_roundtrip_indexname(self, merge_cells, path):
@@ -602,8 +611,8 @@ def test_excel_roundtrip_indexname(self, merge_cells, path):
602611

603612
df.to_excel(path, merge_cells=merge_cells)
604613

605-
xf = ExcelFile(path)
606-
result = pd.read_excel(xf, sheet_name=xf.sheet_names[0], index_col=0)
614+
with ExcelFile(path) as xf:
615+
result = pd.read_excel(xf, sheet_name=xf.sheet_names[0], index_col=0)
607616

608617
tm.assert_frame_equal(result, df)
609618
assert result.index.name == "foo"
@@ -620,8 +629,8 @@ def test_excel_roundtrip_datetime(self, merge_cells, tsframe, path):
620629
tsf.index = [x.date() for x in tsframe.index]
621630
tsf.to_excel(path, "test1", merge_cells=merge_cells)
622631

623-
reader = ExcelFile(path)
624-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
632+
with ExcelFile(path) as reader:
633+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
625634

626635
tm.assert_frame_equal(tsframe, recons)
627636

@@ -680,9 +689,8 @@ def test_to_excel_interval_no_labels(self, path):
680689
expected["new"] = pd.cut(expected[0], 10).astype(str)
681690

682691
df.to_excel(path, "test1")
683-
reader = ExcelFile(path)
684-
685-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
692+
with ExcelFile(path) as reader:
693+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
686694
tm.assert_frame_equal(expected, recons)
687695

688696
def test_to_excel_interval_labels(self, path):
@@ -698,9 +706,8 @@ def test_to_excel_interval_labels(self, path):
698706
expected["new"] = pd.Series(list(intervals))
699707

700708
df.to_excel(path, "test1")
701-
reader = ExcelFile(path)
702-
703-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
709+
with ExcelFile(path) as reader:
710+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
704711
tm.assert_frame_equal(expected, recons)
705712

706713
def test_to_excel_timedelta(self, path):
@@ -718,18 +725,17 @@ def test_to_excel_timedelta(self, path):
718725
)
719726

720727
df.to_excel(path, "test1")
721-
reader = ExcelFile(path)
722-
723-
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
728+
with ExcelFile(path) as reader:
729+
recons = pd.read_excel(reader, sheet_name="test1", index_col=0)
724730
tm.assert_frame_equal(expected, recons)
725731

726732
def test_to_excel_periodindex(self, tsframe, path):
727733
xp = tsframe.resample("M", kind="period").mean()
728734

729735
xp.to_excel(path, "sht1")
730736

731-
reader = ExcelFile(path)
732-
rs = pd.read_excel(reader, sheet_name="sht1", index_col=0)
737+
with ExcelFile(path) as reader:
738+
rs = pd.read_excel(reader, sheet_name="sht1", index_col=0)
733739
tm.assert_frame_equal(xp, rs.to_period("M"))
734740

735741
def test_to_excel_multiindex(self, merge_cells, frame, path):
@@ -742,8 +748,8 @@ def test_to_excel_multiindex(self, merge_cells, frame, path):
742748

743749
# round trip
744750
frame.to_excel(path, "test1", merge_cells=merge_cells)
745-
reader = ExcelFile(path)
746-
df = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
751+
with ExcelFile(path) as reader:
752+
df = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
747753
tm.assert_frame_equal(frame, df)
748754

749755
# GH13511
@@ -771,8 +777,10 @@ def test_to_excel_multiindex_cols(self, merge_cells, frame, path):
771777

772778
# round trip
773779
frame.to_excel(path, "test1", merge_cells=merge_cells)
774-
reader = ExcelFile(path)
775-
df = pd.read_excel(reader, sheet_name="test1", header=header, index_col=[0, 1])
780+
with ExcelFile(path) as reader:
781+
df = pd.read_excel(
782+
reader, sheet_name="test1", header=header, index_col=[0, 1]
783+
)
776784
if not merge_cells:
777785
fm = frame.columns.format(sparsify=False, adjoin=False, names=False)
778786
frame.columns = [".".join(map(str, q)) for q in zip(*fm)]
@@ -785,8 +793,8 @@ def test_to_excel_multiindex_dates(self, merge_cells, tsframe, path):
785793

786794
tsframe.index.names = ["time", "foo"]
787795
tsframe.to_excel(path, "test1", merge_cells=merge_cells)
788-
reader = ExcelFile(path)
789-
recons = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
796+
with ExcelFile(path) as reader:
797+
recons = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1])
790798

791799
tm.assert_frame_equal(tsframe, recons)
792800
assert recons.index.names == ("time", "foo")
@@ -806,8 +814,8 @@ def test_to_excel_multiindex_no_write_index(self, path):
806814
frame2.to_excel(path, "test1", index=False)
807815

808816
# Read it back in.
809-
reader = ExcelFile(path)
810-
frame3 = pd.read_excel(reader, sheet_name="test1")
817+
with ExcelFile(path) as reader:
818+
frame3 = pd.read_excel(reader, sheet_name="test1")
811819

812820
# Test that it is the same as the initial frame.
813821
tm.assert_frame_equal(frame1, frame3)
@@ -820,8 +828,8 @@ def test_to_excel_float_format(self, path):
820828
)
821829
df.to_excel(path, "test1", float_format="%.2f")
822830

823-
reader = ExcelFile(path)
824-
result = pd.read_excel(reader, sheet_name="test1", index_col=0)
831+
with ExcelFile(path) as reader:
832+
result = pd.read_excel(reader, sheet_name="test1", index_col=0)
825833

826834
expected = DataFrame(
827835
[[0.12, 0.23, 0.57], [12.32, 123123.20, 321321.20]],
@@ -979,8 +987,10 @@ def test_excel_010_hemstring(
979987
def roundtrip(data, header=True, parser_hdr=0, index=True):
980988
data.to_excel(path, header=header, merge_cells=merge_cells, index=index)
981989

982-
xf = ExcelFile(path)
983-
return pd.read_excel(xf, sheet_name=xf.sheet_names[0], header=parser_hdr)
990+
with ExcelFile(path) as xf:
991+
return pd.read_excel(
992+
xf, sheet_name=xf.sheet_names[0], header=parser_hdr
993+
)
984994

985995
# Basic test.
986996
parser_header = 0 if use_headers else None

pandas/tests/io/formats/test_info.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,9 @@ def test_info_verbose_with_counts_spacing(
162162
):
163163
"""Test header column, spacer, first line and last line in verbose mode."""
164164
frame = DataFrame(np.random.randn(3, size))
165-
buf = StringIO()
166-
frame.info(verbose=True, show_counts=True, buf=buf)
167-
all_lines = buf.getvalue().splitlines()
165+
with StringIO() as buf:
166+
frame.info(verbose=True, show_counts=True, buf=buf)
167+
all_lines = buf.getvalue().splitlines()
168168
# Here table would contain only header, separator and table lines
169169
# dframe repr, index summary, memory usage and dtypes are excluded
170170
table = all_lines[3:-2]

pandas/tests/io/generate_legacy_storage_files.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -327,9 +327,8 @@ def write_legacy_pickles(output_dir):
327327

328328
pth = f"{platform_name()}.pickle"
329329

330-
fh = open(os.path.join(output_dir, pth), "wb")
331-
pickle.dump(create_pickle_data(), fh, pickle.DEFAULT_PROTOCOL)
332-
fh.close()
330+
with open(os.path.join(output_dir, pth), "wb") as fh:
331+
pickle.dump(create_pickle_data(), fh, pickle.DEFAULT_PROTOCOL)
333332

334333
print(f"created pickle file: {pth}")
335334

pandas/tests/io/parser/test_c_parser_only.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -305,9 +305,9 @@ def test_grow_boundary_at_cap(c_parser_only):
305305
parser = c_parser_only
306306

307307
def test_empty_header_read(count):
308-
s = StringIO("," * count)
309-
expected = DataFrame(columns=[f"Unnamed: {i}" for i in range(count + 1)])
310-
df = parser.read_csv(s)
308+
with StringIO("," * count) as s:
309+
expected = DataFrame(columns=[f"Unnamed: {i}" for i in range(count + 1)])
310+
df = parser.read_csv(s)
311311
tm.assert_frame_equal(df, expected)
312312

313313
for cnt in range(1, 101):

pandas/tests/io/parser/test_network.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,8 @@ def test_read_csv_handles_boto_s3_object(self, s3_resource, tips_file):
250250
Bucket="pandas-test", Key="tips.csv"
251251
)
252252

253-
result = read_csv(BytesIO(s3_object["Body"].read()), encoding="utf8")
253+
with BytesIO(s3_object["Body"].read()) as buffer:
254+
result = read_csv(buffer, encoding="utf8")
254255
assert isinstance(result, DataFrame)
255256
assert not result.empty
256257

0 commit comments

Comments
 (0)