Skip to content

Commit 2e5b05e

Browse files
chrismoradichrismoradi
authored
STYLE: use option_context almost always (#45407)
1 parent ca5d9ee commit 2e5b05e

File tree

9 files changed

+143
-170
lines changed

9 files changed

+143
-170
lines changed

‎pandas/tests/indexes/multi/test_formats.py

+2-6
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
Index,
99
MultiIndex,
1010
)
11-
import pandas._testing as tm
1211

1312

1413
def test_format(idx):
@@ -27,13 +26,10 @@ def test_format_sparse_config(idx):
2726
warn_filters = warnings.filters
2827
warnings.filterwarnings("ignore", category=FutureWarning, module=".*format")
2928
# GH1538
30-
pd.set_option("display.multi_sparse", False)
31-
32-
result = idx.format()
29+
with pd.option_context("display.multi_sparse", False):
30+
result = idx.format()
3331
assert result[1] == "foo two"
3432

35-
tm.reset_display_options()
36-
3733
warnings.filters = warn_filters
3834

3935

‎pandas/tests/indexing/test_chaining_and_caching.py

+10-9
Original file line numberDiff line numberDiff line change
@@ -157,16 +157,17 @@ def test_setitem_chained_setfault(self):
157157
@pytest.mark.arm_slow
158158
def test_detect_chained_assignment(self):
159159

160-
pd.set_option("chained_assignment", "raise")
161-
162-
# work with the chain
163-
expected = DataFrame([[-5, 1], [-6, 3]], columns=list("AB"))
164-
df = DataFrame(np.arange(4).reshape(2, 2), columns=list("AB"), dtype="int64")
165-
assert df._is_copy is None
160+
with option_context("chained_assignment", "raise"):
161+
# work with the chain
162+
expected = DataFrame([[-5, 1], [-6, 3]], columns=list("AB"))
163+
df = DataFrame(
164+
np.arange(4).reshape(2, 2), columns=list("AB"), dtype="int64"
165+
)
166+
assert df._is_copy is None
166167

167-
df["A"][0] = -5
168-
df["A"][1] = -6
169-
tm.assert_frame_equal(df, expected)
168+
df["A"][0] = -5
169+
df["A"][1] = -6
170+
tm.assert_frame_equal(df, expected)
170171

171172
@pytest.mark.arm_slow
172173
def test_detect_chained_assignment_raises(self, using_array_manager):

‎pandas/tests/io/excel/test_writers.py

+4-7
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@
1818
DataFrame,
1919
Index,
2020
MultiIndex,
21-
get_option,
22-
set_option,
21+
option_context,
2322
)
2423
import pandas._testing as tm
2524

@@ -53,10 +52,8 @@ def set_engine(engine, ext):
5352
the test it rolls back said change to the global option.
5453
"""
5554
option_name = f"io.excel.{ext.strip('.')}.writer"
56-
prev_engine = get_option(option_name)
57-
set_option(option_name, engine)
58-
yield
59-
set_option(option_name, prev_engine) # Roll back option change
55+
with option_context(option_name, engine):
56+
yield
6057

6158

6259
@pytest.mark.parametrize(
@@ -1294,7 +1291,7 @@ def check_called(func):
12941291
del called_save[:]
12951292
del called_write_cells[:]
12961293

1297-
with pd.option_context("io.excel.xlsx.writer", "dummy"):
1294+
with option_context("io.excel.xlsx.writer", "dummy"):
12981295
path = "something.xlsx"
12991296
with tm.ensure_clean(path) as filepath:
13001297
register_writer(DummyClass)

‎pandas/tests/io/formats/test_format.py

+20-31
Original file line numberDiff line numberDiff line change
@@ -1146,20 +1146,18 @@ def test_wide_repr(self):
11461146
):
11471147
max_cols = get_option("display.max_columns")
11481148
df = DataFrame(tm.rands_array(25, size=(10, max_cols - 1)))
1149-
set_option("display.expand_frame_repr", False)
1150-
rep_str = repr(df)
1149+
with option_context("display.expand_frame_repr", False):
1150+
rep_str = repr(df)
11511151

11521152
assert f"10 rows x {max_cols - 1} columns" in rep_str
1153-
set_option("display.expand_frame_repr", True)
1154-
wide_repr = repr(df)
1153+
with option_context("display.expand_frame_repr", True):
1154+
wide_repr = repr(df)
11551155
assert rep_str != wide_repr
11561156

11571157
with option_context("display.width", 120):
11581158
wider_repr = repr(df)
11591159
assert len(wider_repr) < len(wide_repr)
11601160

1161-
reset_option("display.expand_frame_repr")
1162-
11631161
def test_wide_repr_wide_columns(self):
11641162
with option_context("mode.sim_interactive", True, "display.max_columns", 20):
11651163
df = DataFrame(
@@ -1174,11 +1172,10 @@ def test_wide_repr_named(self):
11741172
max_cols = get_option("display.max_columns")
11751173
df = DataFrame(tm.rands_array(25, size=(10, max_cols - 1)))
11761174
df.index.name = "DataFrame Index"
1177-
set_option("display.expand_frame_repr", False)
1178-
1179-
rep_str = repr(df)
1180-
set_option("display.expand_frame_repr", True)
1181-
wide_repr = repr(df)
1175+
with option_context("display.expand_frame_repr", False):
1176+
rep_str = repr(df)
1177+
with option_context("display.expand_frame_repr", True):
1178+
wide_repr = repr(df)
11821179
assert rep_str != wide_repr
11831180

11841181
with option_context("display.width", 150):
@@ -1188,18 +1185,16 @@ def test_wide_repr_named(self):
11881185
for line in wide_repr.splitlines()[1::13]:
11891186
assert "DataFrame Index" in line
11901187

1191-
reset_option("display.expand_frame_repr")
1192-
11931188
def test_wide_repr_multiindex(self):
11941189
with option_context("mode.sim_interactive", True, "display.max_columns", 20):
11951190
midx = MultiIndex.from_arrays(tm.rands_array(5, size=(2, 10)))
11961191
max_cols = get_option("display.max_columns")
11971192
df = DataFrame(tm.rands_array(25, size=(10, max_cols - 1)), index=midx)
11981193
df.index.names = ["Level 0", "Level 1"]
1199-
set_option("display.expand_frame_repr", False)
1200-
rep_str = repr(df)
1201-
set_option("display.expand_frame_repr", True)
1202-
wide_repr = repr(df)
1194+
with option_context("display.expand_frame_repr", False):
1195+
rep_str = repr(df)
1196+
with option_context("display.expand_frame_repr", True):
1197+
wide_repr = repr(df)
12031198
assert rep_str != wide_repr
12041199

12051200
with option_context("display.width", 150):
@@ -1209,8 +1204,6 @@ def test_wide_repr_multiindex(self):
12091204
for line in wide_repr.splitlines()[1::13]:
12101205
assert "Level 0 Level 1" in line
12111206

1212-
reset_option("display.expand_frame_repr")
1213-
12141207
def test_wide_repr_multiindex_cols(self):
12151208
with option_context("mode.sim_interactive", True, "display.max_columns", 20):
12161209
max_cols = get_option("display.max_columns")
@@ -1220,34 +1213,30 @@ def test_wide_repr_multiindex_cols(self):
12201213
tm.rands_array(25, (10, max_cols - 1)), index=midx, columns=mcols
12211214
)
12221215
df.index.names = ["Level 0", "Level 1"]
1223-
set_option("display.expand_frame_repr", False)
1224-
rep_str = repr(df)
1225-
set_option("display.expand_frame_repr", True)
1226-
wide_repr = repr(df)
1216+
with option_context("display.expand_frame_repr", False):
1217+
rep_str = repr(df)
1218+
with option_context("display.expand_frame_repr", True):
1219+
wide_repr = repr(df)
12271220
assert rep_str != wide_repr
12281221

12291222
with option_context("display.width", 150, "display.max_columns", 20):
12301223
wider_repr = repr(df)
12311224
assert len(wider_repr) < len(wide_repr)
12321225

1233-
reset_option("display.expand_frame_repr")
1234-
12351226
def test_wide_repr_unicode(self):
12361227
with option_context("mode.sim_interactive", True, "display.max_columns", 20):
12371228
max_cols = 20
12381229
df = DataFrame(tm.rands_array(25, size=(10, max_cols - 1)))
1239-
set_option("display.expand_frame_repr", False)
1240-
rep_str = repr(df)
1241-
set_option("display.expand_frame_repr", True)
1242-
wide_repr = repr(df)
1230+
with option_context("display.expand_frame_repr", False):
1231+
rep_str = repr(df)
1232+
with option_context("display.expand_frame_repr", True):
1233+
wide_repr = repr(df)
12431234
assert rep_str != wide_repr
12441235

12451236
with option_context("display.width", 150):
12461237
wider_repr = repr(df)
12471238
assert len(wider_repr) < len(wide_repr)
12481239

1249-
reset_option("display.expand_frame_repr")
1250-
12511240
def test_wide_repr_wide_long_columns(self):
12521241
with option_context("mode.sim_interactive", True):
12531242
df = DataFrame({"a": ["a" * 30, "b" * 30], "b": ["c" * 70, "d" * 80]})

‎pandas/tests/io/pytables/test_append.py

+54-53
Original file line numberDiff line numberDiff line change
@@ -214,66 +214,66 @@ def test_append_all_nans(setup_path):
214214
tm.assert_frame_equal(store["df2"], df)
215215

216216
# tests the option io.hdf.dropna_table
217-
pd.set_option("io.hdf.dropna_table", False)
218-
_maybe_remove(store, "df3")
219-
store.append("df3", df[:10])
220-
store.append("df3", df[10:])
221-
tm.assert_frame_equal(store["df3"], df)
217+
with pd.option_context("io.hdf.dropna_table", False):
218+
_maybe_remove(store, "df3")
219+
store.append("df3", df[:10])
220+
store.append("df3", df[10:])
221+
tm.assert_frame_equal(store["df3"], df)
222222

223-
pd.set_option("io.hdf.dropna_table", True)
224-
_maybe_remove(store, "df4")
225-
store.append("df4", df[:10])
226-
store.append("df4", df[10:])
227-
tm.assert_frame_equal(store["df4"], df[-4:])
223+
with pd.option_context("io.hdf.dropna_table", True):
224+
_maybe_remove(store, "df4")
225+
store.append("df4", df[:10])
226+
store.append("df4", df[10:])
227+
tm.assert_frame_equal(store["df4"], df[-4:])
228228

229-
# nan some entire rows (string are still written!)
230-
df = DataFrame(
231-
{
232-
"A1": np.random.randn(20),
233-
"A2": np.random.randn(20),
234-
"B": "foo",
235-
"C": "bar",
236-
},
237-
index=np.arange(20),
238-
)
229+
# nan some entire rows (string are still written!)
230+
df = DataFrame(
231+
{
232+
"A1": np.random.randn(20),
233+
"A2": np.random.randn(20),
234+
"B": "foo",
235+
"C": "bar",
236+
},
237+
index=np.arange(20),
238+
)
239239

240-
df.loc[0:15, :] = np.nan
240+
df.loc[0:15, :] = np.nan
241241

242-
_maybe_remove(store, "df")
243-
store.append("df", df[:10], dropna=True)
244-
store.append("df", df[10:], dropna=True)
245-
tm.assert_frame_equal(store["df"], df)
242+
_maybe_remove(store, "df")
243+
store.append("df", df[:10], dropna=True)
244+
store.append("df", df[10:], dropna=True)
245+
tm.assert_frame_equal(store["df"], df)
246246

247-
_maybe_remove(store, "df2")
248-
store.append("df2", df[:10], dropna=False)
249-
store.append("df2", df[10:], dropna=False)
250-
tm.assert_frame_equal(store["df2"], df)
247+
_maybe_remove(store, "df2")
248+
store.append("df2", df[:10], dropna=False)
249+
store.append("df2", df[10:], dropna=False)
250+
tm.assert_frame_equal(store["df2"], df)
251251

252-
# nan some entire rows (but since we have dates they are still
253-
# written!)
254-
df = DataFrame(
255-
{
256-
"A1": np.random.randn(20),
257-
"A2": np.random.randn(20),
258-
"B": "foo",
259-
"C": "bar",
260-
"D": Timestamp("20010101"),
261-
"E": datetime.datetime(2001, 1, 2, 0, 0),
262-
},
263-
index=np.arange(20),
264-
)
252+
# nan some entire rows (but since we have dates they are still
253+
# written!)
254+
df = DataFrame(
255+
{
256+
"A1": np.random.randn(20),
257+
"A2": np.random.randn(20),
258+
"B": "foo",
259+
"C": "bar",
260+
"D": Timestamp("20010101"),
261+
"E": datetime.datetime(2001, 1, 2, 0, 0),
262+
},
263+
index=np.arange(20),
264+
)
265265

266-
df.loc[0:15, :] = np.nan
266+
df.loc[0:15, :] = np.nan
267267

268-
_maybe_remove(store, "df")
269-
store.append("df", df[:10], dropna=True)
270-
store.append("df", df[10:], dropna=True)
271-
tm.assert_frame_equal(store["df"], df)
268+
_maybe_remove(store, "df")
269+
store.append("df", df[:10], dropna=True)
270+
store.append("df", df[10:], dropna=True)
271+
tm.assert_frame_equal(store["df"], df)
272272

273-
_maybe_remove(store, "df2")
274-
store.append("df2", df[:10], dropna=False)
275-
store.append("df2", df[10:], dropna=False)
276-
tm.assert_frame_equal(store["df2"], df)
273+
_maybe_remove(store, "df2")
274+
store.append("df2", df[:10], dropna=False)
275+
store.append("df2", df[10:], dropna=False)
276+
tm.assert_frame_equal(store["df2"], df)
277277

278278

279279
def test_append_frame_column_oriented(setup_path):
@@ -898,8 +898,9 @@ def test_append_to_multiple_dropna_false(setup_path):
898898
df1.iloc[1, df1.columns.get_indexer(["A", "B"])] = np.nan
899899
df = concat([df1, df2], axis=1)
900900

901-
with ensure_clean_store(setup_path) as store:
902-
901+
with ensure_clean_store(setup_path) as store, pd.option_context(
902+
"io.hdf.dropna_table", True
903+
):
903904
# dropna=False shouldn't synchronize row indexes
904905
store.append_to_multiple(
905906
{"df1a": ["A", "B"], "df2a": None}, df, selector="df1a", dropna=False

0 commit comments

Comments
 (0)