BUG: support partial setting with .ix in Series and DataFrame, refactoring, GH #397

Author: wesm

TODO.rst

@@ -7,6 +7,7 @@ TODO
 ----
 - _consolidate, does it always copy?
 - Series.align with fill method. Will have to generate more Cython code
+- TYPE inference in Index-- more than just datetime!
 
 TODO docs
 ---------

pandas/core/indexing.py

@@ -36,21 +36,33 @@ def _getitem_xs(self, idx, axis=0):
             return self.obj.xs(idx, axis=axis, copy=True)
 
     def __setitem__(self, key, value):
+        # kludgetastic
+        ax = self.obj._get_axis(0)
+        if isinstance(ax, MultiIndex):
+            try:
+                indexer = ax.get_loc(key)
+                self._setitem_with_indexer(indexer, value)
+                return
+            except Exception:
+                pass
+
         if isinstance(key, tuple):
             if len(key) > self.ndim:
                 raise IndexingError('only tuples of length <= %d supported',
                                     self.ndim)
-
-            keyidx = []
-            for i, k in enumerate(key):
-                idx = self._convert_to_indexer(k, axis=i)
-                keyidx.append(idx)
-            indexer = _maybe_convert_ix(*keyidx)
+            indexer = self._convert_tuple(key)
         else:
             indexer = self._convert_to_indexer(key)
 
         self._setitem_with_indexer(indexer, value)
 
+    def _convert_tuple(self, key):
+        keyidx = []
+        for i, k in enumerate(key):
+            idx = self._convert_to_indexer(k, axis=i)
+            keyidx.append(idx)
+        return _maybe_convert_ix(*keyidx)
+
     def _setitem_with_indexer(self, indexer, value):
         # also has the side effect of consolidating in-place
         if self.obj._is_mixed_type:
@@ -274,51 +286,73 @@ class _SeriesIndexer(_NDFrameIndexer):
     """
 
     def __getitem__(self, key):
-        op = self._fancy_index(key, operation='get')
-        return op()
-
-    def __setitem__(self, key, value):
-        op = self._fancy_index(key, value, operation='set')
-        op()
-
-    def _fancy_index(self, key, value=None, operation='get'):
-        # going to great lengths to avoid code dup
-        obj = self.obj
+        ax = self.obj.index
+        if isinstance(ax, MultiIndex):
+            try:
+                # key = ax.get_loc(key)
+                return self._get_default(key)
+            except Exception:
+                pass
 
-        if operation == 'get':
-            def do_default():
-                return obj[key]
+        if _isboolarr(key):
+            self._check_boolean_key(key)
+        elif isinstance(key, slice):
+            key = self._convert_slice(key)
+        elif _is_list_like(key):
+            return self._get_list_like(key)
+        return self._get_default(key)
 
-            def do_list_like():
-                if isinstance(obj.index, MultiIndex):
-                    try:
-                        return obj[key]
-                    except (KeyError, TypeError, IndexError):
-                        pass
-                return obj.reindex(key)
-        else:
-            def do_default():
-                obj[key] = value
+    def __setitem__(self, key, value):
+        ax = self.obj.index
+        if isinstance(ax, MultiIndex):
+            try:
+                key = ax.get_loc(key)
+                self._set_default(key, value)
+                return
+            except Exception:
+                pass
 
-            def do_list_like():
-                inds = obj.index.get_indexer(key)
-                mask = inds == -1
-                if mask.any():
-                    raise IndexingError('Indices %s not found' % key[mask])
-                obj.put(inds, value)
-        op = do_default
         if _isboolarr(key):
-            if _is_series(key):
-                if not key.index.equals(obj.index):
-                    raise IndexingError('Cannot use boolean index with '
-                                        'misaligned or unequal labels')
+            self._check_boolean_key(key)
         elif isinstance(key, slice):
-            if _is_label_slice(obj.index, key):
-                i, j = obj.index.slice_locs(key.start, key.stop)
-                key = slice(i, j)
+            key = self._convert_slice(key)
         elif _is_list_like(key):
-            op = do_list_like
-        return op
+            self._set_list_like(key, value)
+            return
+        return self._set_default(key, value)
+
+    def _check_boolean_key(self, key):
+        if _is_series(key):
+            if not key.index.equals(self.obj.index):
+                raise IndexingError('Cannot use boolean index with '
+                                    'misaligned or unequal labels')
+
+    def _convert_slice(self, key):
+        if _is_label_slice(self.obj.index, key):
+            i, j = self.obj.index.slice_locs(key.start, key.stop)
+            key = slice(i, j)
+        return key
+
+    def _get_default(self, key):
+        return self.obj[key]
+
+    def _get_list_like(self, key):
+        if isinstance(self.obj.index, MultiIndex):
+            try:
+                return self.obj[key]
+            except (KeyError, TypeError, IndexError):
+                pass
+        return self.obj.reindex(key)
+
+    def _set_default(self, key, value):
+        self.obj[key] = value
+
+    def _set_list_like(self, key, value):
+        inds = self.obj.index.get_indexer(key)
+        mask = inds == -1
+        if mask.any():
+            raise IndexingError('Indices %s not found' % key[mask])
+        self.obj.put(inds, value)
 
 def _is_series(obj):
     from pandas.core.series import Series

pandas/tests/test_multilevel.py

@@ -663,6 +663,22 @@ def test_ix_preserve_names(self):
         self.assertEquals(result.index.name, self.ymd.index.names[2])
         self.assertEquals(result2.index.name, self.ymd.index.names[2])
 
+    def test_partial_set(self):
+        # GH #397
+        df = self.ymd.copy()
+        exp = self.ymd.copy()
+        df.ix[2000, 4] = 0
+        exp.ix[2000, 4].values[:] = 0
+        assert_frame_equal(df, exp)
+
+        df['A'].ix[2000, 4] = 1
+        exp['A'].ix[2000, 4].values[:] = 1
+        assert_frame_equal(df, exp)
+
+        df.ix[2000] = 5
+        exp.ix[2000].values[:] = 5
+        assert_frame_equal(df, exp)
+
 if __name__ == '__main__':
 
     # unittest.main()