Refactor box constructors

pandas-dev · Aug 9, 2018 · 7aa1af4 · 7aa1af4
1 parent ee569b3
commit 7aa1af4
Showing 1 changed file with 23 additions and 9 deletions.
diff --git a/pandas/tests/frame/test_alter_axes.py b/pandas/tests/frame/test_alter_axes.py
@@ -22,10 +22,6 @@
 
 from pandas.tests.frame.common import TestData
 
-key = lambda x: x.name
-mi = lambda x: MultiIndex.from_arrays([x])
-
-
 class TestDataFrameAlterAxes(TestData):
 
     def test_set_index_directly(self):
@@ -116,14 +112,18 @@ def test_set_index_after_mutation(self):
         tm.assert_frame_equal(result, expected)
 
     # also test index name if append=True (name is duplicate here for B)
-    @pytest.mark.parametrize('box', [Series, Index, np.array, mi])
+    @pytest.mark.parametrize('box', [Series, Index, np.array, 'MultiIndex'])
     @pytest.mark.parametrize('append, index_name', [(True, None),
                              (True, 'B'), (True, 'test'), (False, None)])
     @pytest.mark.parametrize('drop', [True, False])
     def test_set_index_pass_single_array(self, drop, append, index_name, box):
         df = self.dummy.copy()
         df.index.name = index_name
 
+        # update constructor in case of MultiIndex
+        box = ((lambda x: MultiIndex.from_arrays([x]))
+               if box == 'MultiIndex' else box)
+
         key = box(df['B'])
         # np.array and list "forget" the name of B
         name = [None if box in [np.array, list] else 'B']
@@ -138,7 +138,7 @@ def test_set_index_pass_single_array(self, drop, append, index_name, box):
         tm.assert_frame_equal(result, expected)
 
     # also test index name if append=True (name is duplicate here for A & B)
-    @pytest.mark.parametrize('box', [Series, Index, np.array, list, mi])
+    @pytest.mark.parametrize('box', [Series, Index, np.array, list, 'MultiIndex'])
     @pytest.mark.parametrize('append, index_name',
                              [(True, None), (True, 'A'), (True, 'B'),
                               (True, 'test'), (False, None)])
@@ -147,6 +147,10 @@ def test_set_index_pass_arrays(self, drop, append, index_name, box):
         df = self.dummy.copy()
         df.index.name = index_name
 
+        # update constructor in case of MultiIndex
+        box = ((lambda x: MultiIndex.from_arrays([x]))
+               if box == 'MultiIndex' else box)
+
         keys = ['A', box(df['B'])]
         # np.array and list "forget" the name of B
         names = ['A', None if box in [np.array, list] else 'B']
@@ -162,8 +166,10 @@ def test_set_index_pass_arrays(self, drop, append, index_name, box):
         tm.assert_frame_equal(result, expected)
 
     # also test index name if append=True (name is duplicate here for A)
-    @pytest.mark.parametrize('box1', [key, Series, Index, np.array, list, mi])
-    @pytest.mark.parametrize('box2', [key, Series, Index, np.array, list, mi])
+    @pytest.mark.parametrize('box1', ['label', Series, Index, np.array,
+                                      list, 'MultiIndex'])
+    @pytest.mark.parametrize('box2', ['label', Series, Index, np.array,
+                                      list, 'MultiIndex'])
     @pytest.mark.parametrize('append, index_name', [(True, None),
                              (True, 'A'), (True, 'test'), (False, None)])
     @pytest.mark.parametrize('drop', [True, False])
@@ -172,7 +178,15 @@ def test_set_index_pass_arrays_duplicate(self, drop, append, index_name,
         df = self.dummy.copy()
         df.index.name = index_name
 
-        keys = [box1(df['A']), box2(df['A'])]
+        # transform strings to correct box constructor
+        def rebox(x):
+            if x == 'label':
+                return lambda x: x.name
+            elif x == 'MultiIndex':
+                return lambda x: MultiIndex.from_arrays([x])
+            return x
+
+        keys = [rebox(box1)(df['A']), rebox(box2)(df['A'])]
 
         # == gives ambiguous Boolean for Series
         if keys[0] is 'A' and keys[1] is 'A':