Skip to content

Commit 8df5acf

Browse files
committed
modified: pandas/tests/reshape/test_pivot_multilevel.py
- Combined tests :func:`test_pivot_table_values_in_columns` and :func:`test_pivot_table_values_in_index` into a single parametrized test, :func:`test_pivot_table_multiindex_values_as_two_params` to reduce duplicate setup code.
1 parent a1f7694 commit 8df5acf

File tree

1 file changed

+29
-33
lines changed

1 file changed

+29
-33
lines changed

pandas/tests/reshape/test_pivot_multilevel.py

+29-33
Original file line numberDiff line numberDiff line change
@@ -252,47 +252,43 @@ def test_pivot_df_multiindex_index_none():
252252
tm.assert_frame_equal(result, expected)
253253

254254

255-
def test_pivot_table_values_in_columns():
255+
@pytest.mark.parametrize(
256+
argnames=["index", "columns"],
257+
argvalues=[("index", ["col", "value"]), (["index", "value"], "col")],
258+
ids=["values-and-columns", "values-and-index"],
259+
)
260+
def test_pivot_table_multiindex_values_as_two_params(index, columns, request):
256261
data = [
257262
["A", 1, 50, -1],
258263
["B", 1, 100, -2],
259264
["A", 2, 100, -2],
260265
["B", 2, 200, -4],
261266
]
262267
df = pd.DataFrame(data=data, columns=["index", "col", "value", "extra"])
263-
result = df.pivot_table(values="value", index="index", columns=["col", "value"])
268+
result = df.pivot_table(values="value", index=index, columns=columns)
264269
nan = np.nan
265-
e_data = [
266-
[50.0, nan, 100.0, nan],
267-
[nan, 100.0, nan, 200.0],
268-
]
269-
e_index = Index(data=["A", "B"], name="index")
270-
e_cols = MultiIndex.from_arrays(
271-
arrays=[[1, 1, 2, 2], [50, 100, 100, 200]], names=["col", "value"]
272-
)
273-
expected = pd.DataFrame(data=e_data, index=e_index, columns=e_cols)
274-
tm.assert_frame_equal(left=result, right=expected)
270+
if request.node.callspec.id == "values-and-columns":
271+
e_data = [
272+
[50.0, nan, 100.0, nan],
273+
[nan, 100.0, nan, 200.0],
274+
]
275+
e_index = Index(data=["A", "B"], name="index")
276+
e_cols = MultiIndex.from_arrays(
277+
arrays=[[1, 1, 2, 2], [50, 100, 100, 200]], names=["col", "value"]
278+
)
279+
expected = pd.DataFrame(data=e_data, index=e_index, columns=e_cols)
275280

281+
else:
282+
e_data = [
283+
[50.0, nan],
284+
[nan, 100.0],
285+
[100.0, nan],
286+
[nan, 200.0],
287+
]
288+
e_index = MultiIndex.from_arrays(
289+
arrays=[["A", "A", "B", "B"], [50, 100, 100, 200]], names=["index", "value"]
290+
)
291+
e_cols = Index(data=[1, 2], name="col")
292+
expected = pd.DataFrame(data=e_data, index=e_index, columns=e_cols)
276293

277-
def test_pivot_table_values_in_index():
278-
data = [
279-
["A", 1, 50, -1],
280-
["B", 1, 100, -2],
281-
["A", 2, 100, -2],
282-
["B", 2, 200, -4],
283-
]
284-
df = pd.DataFrame(data=data, columns=["index", "col", "value", "extra"])
285-
result = df.pivot_table(values="value", index=["index", "value"], columns="col")
286-
nan = np.nan
287-
e_data = [
288-
[50.0, nan],
289-
[nan, 100.0],
290-
[100.0, nan],
291-
[nan, 200.0],
292-
]
293-
e_index = MultiIndex.from_arrays(
294-
arrays=[["A", "A", "B", "B"], [50, 100, 100, 200]], names=["index", "value"]
295-
)
296-
e_cols = Index(data=[1, 2], name="col")
297-
expected = pd.DataFrame(data=e_data, index=e_index, columns=e_cols)
298294
tm.assert_frame_equal(left=result, right=expected)

0 commit comments

Comments
 (0)