ENH: partial multiple setting on first level via .ix on DataFrame, close #409

Author: wesm

pandas/core/indexing.py

@@ -197,7 +197,6 @@ def _getitem_axis(self, key, axis=0):
 
     def _getitem_iterable(self, key, axis=0):
         labels = self.obj._get_axis(axis)
-        # axis_name = self.obj._get_axis_name(axis)
 
         def _reindex(keys, level=None):
             try:
@@ -305,8 +304,18 @@ def _convert_to_indexer(self, obj, axis=0):
                 if _is_integer_dtype(objarr) and not is_int_index:
                     return objarr
 
-                indexer = labels.get_indexer(objarr)
-                mask = indexer == -1
+                # this is not the most robust, but...
+                if (isinstance(labels, MultiIndex) and
+                    not isinstance(objarr[0], tuple)):
+                    level = 0
+                    _, indexer = labels.reindex(objarr, level=level)
+
+                    check = labels.levels[0].get_indexer(objarr)
+                else:
+                    level = None
+                    indexer = check = labels.get_indexer(objarr)
+
+                mask = check == -1
                 if mask.any():
                     raise KeyError('%s not in index' % objarr[mask])
 

pandas/tests/test_multilevel.py

@@ -1233,6 +1233,35 @@ def test_reindex_level_partial_selection(self):
         result = self.frame.T.ix[:, ['foo', 'qux']]
         assert_frame_equal(result, expected.T)
 
+    def test_setitem_multiple_partial(self):
+        expected = self.frame.copy()
+        result = self.frame.copy()
+        result.ix[['foo', 'bar']] = 0
+        expected.ix['foo'] = 0
+        expected.ix['bar'] = 0
+        assert_frame_equal(result, expected)
+
+        expected = self.frame.copy()
+        result = self.frame.copy()
+        result.ix['foo':'bar'] = 0
+        expected.ix['foo'] = 0
+        expected.ix['bar'] = 0
+        assert_frame_equal(result, expected)
+
+        expected = self.frame['A'].copy()
+        result = self.frame['A'].copy()
+        result.ix[['foo', 'bar']] = 0
+        expected.ix['foo'] = 0
+        expected.ix['bar'] = 0
+        assert_series_equal(result, expected)
+
+        expected = self.frame['A'].copy()
+        result = self.frame['A'].copy()
+        result.ix['foo':'bar'] = 0
+        expected.ix['foo'] = 0
+        expected.ix['bar'] = 0
+        assert_series_equal(result, expected)
+
     def test_drop_level(self):
         result = self.frame.drop(['bar', 'qux'], level='first')
         expected = self.frame.ix[[0, 1, 2, 5, 6]]