22
22
23
23
from pandas .tests .frame .common import TestData
24
24
25
- key = lambda x : x .name
26
- mi = lambda x : MultiIndex .from_arrays ([x ])
27
-
28
25
29
26
class TestDataFrameAlterAxes (TestData ):
30
27
@@ -116,14 +113,18 @@ def test_set_index_after_mutation(self):
116
113
tm .assert_frame_equal (result , expected )
117
114
118
115
# 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' ])
120
117
@pytest .mark .parametrize ('append, index_name' , [(True , None ),
121
118
(True , 'B' ), (True , 'test' ), (False , None )])
122
119
@pytest .mark .parametrize ('drop' , [True , False ])
123
120
def test_set_index_pass_single_array (self , drop , append , index_name , box ):
124
121
df = self .dummy .copy ()
125
122
df .index .name = index_name
126
123
124
+ # update constructor in case of MultiIndex
125
+ box = ((lambda x : MultiIndex .from_arrays ([x ]))
126
+ if box == 'MultiIndex' else box )
127
+
127
128
key = box (df ['B' ])
128
129
# np.array and list "forget" the name of B
129
130
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):
138
139
tm .assert_frame_equal (result , expected )
139
140
140
141
# 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' ])
142
144
@pytest .mark .parametrize ('append, index_name' ,
143
145
[(True , None ), (True , 'A' ), (True , 'B' ),
144
146
(True , 'test' ), (False , None )])
@@ -147,6 +149,10 @@ def test_set_index_pass_arrays(self, drop, append, index_name, box):
147
149
df = self .dummy .copy ()
148
150
df .index .name = index_name
149
151
152
+ # update constructor in case of MultiIndex
153
+ box = ((lambda x : MultiIndex .from_arrays ([x ]))
154
+ if box == 'MultiIndex' else box )
155
+
150
156
keys = ['A' , box (df ['B' ])]
151
157
# np.array and list "forget" the name of B
152
158
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):
162
168
tm .assert_frame_equal (result , expected )
163
169
164
170
# 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' ])
167
175
@pytest .mark .parametrize ('append, index_name' , [(True , None ),
168
176
(True , 'A' ), (True , 'test' ), (False , None )])
169
177
@pytest .mark .parametrize ('drop' , [True , False ])
@@ -172,7 +180,15 @@ def test_set_index_pass_arrays_duplicate(self, drop, append, index_name,
172
180
df = self .dummy .copy ()
173
181
df .index .name = index_name
174
182
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' ])]
176
192
177
193
# == gives ambiguous Boolean for Series
178
194
if keys [0 ] is 'A' and keys [1 ] is 'A' :
0 commit comments