Skip to content

Commit f46c1e4

Browse files
committed
Refactor box constructors
1 parent ee569b3 commit f46c1e4

File tree

1 file changed

+24
-8
lines changed

1 file changed

+24
-8
lines changed

pandas/tests/frame/test_alter_axes.py

+24-8
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,6 @@
2222

2323
from pandas.tests.frame.common import TestData
2424

25-
key = lambda x: x.name
26-
mi = lambda x: MultiIndex.from_arrays([x])
27-
2825

2926
class TestDataFrameAlterAxes(TestData):
3027

@@ -116,14 +113,18 @@ def test_set_index_after_mutation(self):
116113
tm.assert_frame_equal(result, expected)
117114

118115
# also test index name if append=True (name is duplicate here for B)
119-
@pytest.mark.parametrize('box', [Series, Index, np.array, mi])
116+
@pytest.mark.parametrize('box', [Series, Index, np.array, 'MultiIndex'])
120117
@pytest.mark.parametrize('append, index_name', [(True, None),
121118
(True, 'B'), (True, 'test'), (False, None)])
122119
@pytest.mark.parametrize('drop', [True, False])
123120
def test_set_index_pass_single_array(self, drop, append, index_name, box):
124121
df = self.dummy.copy()
125122
df.index.name = index_name
126123

124+
# update constructor in case of MultiIndex
125+
box = ((lambda x: MultiIndex.from_arrays([x]))
126+
if box == 'MultiIndex' else box)
127+
127128
key = box(df['B'])
128129
# np.array and list "forget" the name of B
129130
name = [None if box in [np.array, list] else 'B']
@@ -138,7 +139,8 @@ def test_set_index_pass_single_array(self, drop, append, index_name, box):
138139
tm.assert_frame_equal(result, expected)
139140

140141
# also test index name if append=True (name is duplicate here for A & B)
141-
@pytest.mark.parametrize('box', [Series, Index, np.array, list, mi])
142+
@pytest.mark.parametrize('box', [Series, Index, np.array,
143+
list, 'MultiIndex'])
142144
@pytest.mark.parametrize('append, index_name',
143145
[(True, None), (True, 'A'), (True, 'B'),
144146
(True, 'test'), (False, None)])
@@ -147,6 +149,10 @@ def test_set_index_pass_arrays(self, drop, append, index_name, box):
147149
df = self.dummy.copy()
148150
df.index.name = index_name
149151

152+
# update constructor in case of MultiIndex
153+
box = ((lambda x: MultiIndex.from_arrays([x]))
154+
if box == 'MultiIndex' else box)
155+
150156
keys = ['A', box(df['B'])]
151157
# np.array and list "forget" the name of B
152158
names = ['A', None if box in [np.array, list] else 'B']
@@ -162,8 +168,10 @@ def test_set_index_pass_arrays(self, drop, append, index_name, box):
162168
tm.assert_frame_equal(result, expected)
163169

164170
# also test index name if append=True (name is duplicate here for A)
165-
@pytest.mark.parametrize('box1', [key, Series, Index, np.array, list, mi])
166-
@pytest.mark.parametrize('box2', [key, Series, Index, np.array, list, mi])
171+
@pytest.mark.parametrize('box1', ['label', Series, Index, np.array,
172+
list, 'MultiIndex'])
173+
@pytest.mark.parametrize('box2', ['label', Series, Index, np.array,
174+
list, 'MultiIndex'])
167175
@pytest.mark.parametrize('append, index_name', [(True, None),
168176
(True, 'A'), (True, 'test'), (False, None)])
169177
@pytest.mark.parametrize('drop', [True, False])
@@ -172,7 +180,15 @@ def test_set_index_pass_arrays_duplicate(self, drop, append, index_name,
172180
df = self.dummy.copy()
173181
df.index.name = index_name
174182

175-
keys = [box1(df['A']), box2(df['A'])]
183+
# transform strings to correct box constructor
184+
def rebox(x):
185+
if x == 'label':
186+
return lambda x: x.name
187+
elif x == 'MultiIndex':
188+
return lambda x: MultiIndex.from_arrays([x])
189+
return x
190+
191+
keys = [rebox(box1)(df['A']), rebox(box2)(df['A'])]
176192

177193
# == gives ambiguous Boolean for Series
178194
if keys[0] is 'A' and keys[1] is 'A':

0 commit comments

Comments
 (0)