Tests for 'mixed depth' hierarchical columns.

Tests for getting, setting, dropping and popping columns from dataframes which have MultiIndexed columns containing some empty strings at the lower levels, by indexing with only the upper levels.
Also modified NDFrame._set_item to use the same shorthand.

Author: chrisjbillington

pandas/core/generic.py

@@ -291,6 +291,13 @@ def _clear_item_cache(self):
         self._item_cache.clear()
 
     def _set_item(self, key, value):
+        if isinstance(self.columns, MultiIndex):
+            # Pad the key with empty strings if lower levels of the key
+            # aren't specified:
+            if not isinstance(key, tuple):
+                key = (key,)
+            if len(key) != self.columns.nlevels:
+                key += ('',)*(self.columns.nlevels - len(key))
         self._data.set(key, value)
 
         try:
@@ -303,9 +310,9 @@ def __delitem__(self, key):
         Delete item
         """
         deleted = False
-        if key not in self.columns:
-            # If column labels are tuples, allow shorthand to delete
-            # all columns whose first len(key) elements match key:
+        if isinstance(self.columns, MultiIndex) and key not in self.columns:
+            # Allow shorthand to delete all columns whose first len(key)
+            # elements match key:
             if not isinstance(key,tuple):
                 key = (key,)
             for col in self.columns:

pandas/tests/test_multilevel.py

@@ -1083,6 +1083,86 @@ def test_int_series_slicing(self):
         expected = self.ymd.reindex(s.index[5:])
         assert_frame_equal(result, expected)
 
+    def test_mixed_depth_get(self):
+        arrays = [[  'a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
+                  [   '',  'OD',  'OD', 'result1',   'result2',  'result1'],
+                  [   '',  'wx',  'wy',        '',          '',         '']]
+
+        tuples = zip(*arrays)
+        tuples.sort()
+        index = MultiIndex.from_tuples(tuples)
+        df = DataFrame(randn(4,6),columns = index)
+            
+        result = df['a']
+        expected = df['a','','']
+        assert_series_equal(result, expected)
+        self.assertEquals(result.name, 'a')
+        
+        result = df['routine1','result1']
+        expected = df['routine1','result1','']
+        assert_series_equal(result, expected)
+        self.assertEquals(result.name, ('routine1', 'result1'))
+
+    def test_mixed_depth_insert(self):  
+        arrays = [[  'a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
+                  [   '',  'OD',  'OD', 'result1',   'result2',  'result1'],
+                  [   '',  'wx',  'wy',        '',          '',         '']]
+
+        tuples = zip(*arrays)
+        tuples.sort()
+        index = MultiIndex.from_tuples(tuples)
+        df = DataFrame(randn(4,6),columns = index)     
+
+        result = df.copy()
+        expected = df.copy()
+        result['b'] = [1,2,3,4]
+        expected['b','',''] = [1,2,3,4]
+        assert_frame_equal(result, expected)
+ 
+    def test_mixed_depth_drop(self):  
+        arrays = [[  'a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
+                  [   '',  'OD',  'OD', 'result1',   'result2',  'result1'],
+                  [   '',  'wx',  'wy',        '',          '',         '']]
+
+        tuples = zip(*arrays)
+        tuples.sort()
+        index = MultiIndex.from_tuples(tuples)
+        df = DataFrame(randn(4,6),columns = index)     
+
+        result = df.drop('a',axis=1)
+        expected = df.drop([('a','','')],axis=1)
+        assert_frame_equal(expected, result)
+        
+        result = df.drop(['top'],axis=1)
+        expected = df.drop([('top','OD','wx')], axis=1)
+        expected = expected.drop([('top','OD','wy')], axis=1)
+        assert_frame_equal(expected, result)
+               
+    def test_mixed_depth_pop(self):  
+        arrays = [[  'a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
+                  [   '',  'OD',  'OD', 'result1',   'result2',  'result1'],
+                  [   '',  'wx',  'wy',        '',          '',         '']]
+
+        tuples = zip(*arrays)
+        tuples.sort()
+        index = MultiIndex.from_tuples(tuples)
+        df = DataFrame(randn(4,6),columns = index)     
+
+        df1 = df.copy()
+        df2 = df.copy()
+        result = df1.pop('a')
+        expected = df2.pop(('a','',''))
+        assert_series_equal(expected, result)
+        assert_frame_equal(df1, df2)
+        self.assertEquals(result.name,'a')
+        
+        expected = df1['top']
+        df1 = df1.drop(['top'],axis=1)
+        result = df2.pop('top')
+        assert_frame_equal(expected, result)
+        assert_frame_equal(df1, df2)
+            
+        
 if __name__ == '__main__':
 
     # unittest.main()