ENH: Add set_index to Series (WIP)

Author: h-vetinari

pandas/core/frame.py

@@ -3928,6 +3928,10 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
             necessary. Setting to False will improve the performance of this
             method
 
+        Returns
+        -------
+        reindexed : DataFrame if inplace is False, else None
+
         Examples
         --------
         >>> df = pd.DataFrame({'month': [1, 4, 7, 10],
@@ -3968,73 +3972,11 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
         2  2014  4      40
         3  2013  7      84
         4  2014  10     31
-
-        Returns
-        -------
-        dataframe : DataFrame
         """
-        inplace = validate_bool_kwarg(inplace, 'inplace')
-        if not isinstance(keys, list):
-            keys = [keys]
-
-        if inplace:
-            frame = self
-        else:
-            frame = self.copy()
-
-        arrays = []
-        names = []
-        if append:
-            names = [x for x in self.index.names]
-            if isinstance(self.index, MultiIndex):
-                for i in range(self.index.nlevels):
-                    arrays.append(self.index._get_level_values(i))
-            else:
-                arrays.append(self.index)
-
-        to_remove = []
-        for col in keys:
-            if isinstance(col, MultiIndex):
-                # append all but the last column so we don't have to modify
-                # the end of this loop
-                for n in range(col.nlevels - 1):
-                    arrays.append(col._get_level_values(n))
-
-                level = col._get_level_values(col.nlevels - 1)
-                names.extend(col.names)
-            elif isinstance(col, Series):
-                level = col._values
-                names.append(col.name)
-            elif isinstance(col, Index):
-                level = col
-                names.append(col.name)
-            elif isinstance(col, (list, np.ndarray, Index)):
-                level = col
-                names.append(None)
-            else:
-                level = frame[col]._values
-                names.append(col)
-                if drop:
-                    to_remove.append(col)
-            arrays.append(level)
-
-        index = ensure_index_from_sequences(arrays, names)
-
-        if verify_integrity and not index.is_unique:
-            duplicates = index[index.duplicated()].unique()
-            raise ValueError('Index has duplicate keys: {dup}'.format(
-                dup=duplicates))
-
-        for c in to_remove:
-            del frame[c]
-
-        # clear up memory usage
-        index._cleanup()
-
-        frame.index = index
-
-        if not inplace:
-            return frame
+        vi = verify_integrity
+        return super(DataFrame, self).set_index(keys=keys, drop=drop,
+                                                append=append, inplace=inplace,
+                                                verify_integrity=vi)
 
     def reset_index(self, level=None, drop=False, inplace=False, col_level=0,
                     col_fill=''):

pandas/core/generic.py

@@ -35,8 +35,9 @@
 from pandas.core.dtypes.generic import ABCSeries, ABCPanel, ABCDataFrame
 
 from pandas.core.base import PandasObject, SelectionMixin
-from pandas.core.index import (Index, MultiIndex, ensure_index,
-                               InvalidIndexError, RangeIndex)
+from pandas.core.index import (Index, MultiIndex,
+                               InvalidIndexError, RangeIndex,
+                               ensure_index, ensure_index_from_sequences)
 import pandas.core.indexing as indexing
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.core.indexes.period import PeriodIndex, Period
@@ -666,6 +667,141 @@ def _set_axis(self, axis, labels):
         y : same as input
         """
 
+    def set_index(self, keys, drop=True, append=False, inplace=False,
+                  verify_integrity=False):
+        """
+        Set the Series/DataFrame index (row labels) using one or more given
+        arrays (or column labels in case of DataFrame).
+        By default yields a new object.
+
+        Parameters
+        ----------
+        keys : column label or list of column labels / arrays. For Series case,
+            only array or list of arrays is allowed.
+        drop : boolean, default True
+            Delete columns to be used as the new index (only for DataFrame).
+        append : boolean, default False
+            Whether to append columns to existing index
+        inplace : boolean, default False
+            Modify the Series/DataFrame in place (do not create a new object)
+        verify_integrity : boolean, default False
+            Check the new index for duplicates. Otherwise defer the check until
+            necessary. Setting to False will improve the performance of this
+            method
+
+        Returns
+        -------
+        reindexed : Series/DataFrame if inplace is False, else None
+
+        Examples
+        --------
+        >>> df = pd.DataFrame({'month': [1, 4, 7, 10],
+        ...                    'year': [2012, 2014, 2013, 2014],
+        ...                    'sale':[55, 40, 84, 31]})
+           month  sale  year
+        0  1      55    2012
+        1  4      40    2014
+        2  7      84    2013
+        3  10     31    2014
+
+        Set the index to become the 'month' column:
+
+        >>> df.set_index('month')
+               sale  year
+        month
+        1      55    2012
+        4      40    2014
+        7      84    2013
+        10     31    2014
+
+        Create a multi-index using columns 'year' and 'month':
+
+        >>> df.set_index(['year', 'month'])
+                    sale
+        year  month
+        2012  1     55
+        2014  4     40
+        2013  7     84
+        2014  10    31
+
+        Create a multi-index using a set of values and a column:
+
+        >>> df.set_index([[1, 2, 3, 4], 'year'])
+                 month  sale
+           year
+        1  2012  1      55
+        2  2014  4      40
+        3  2013  7      84
+        4  2014  10     31
+        """
+        from pandas import Series
+
+        inplace = validate_bool_kwarg(inplace, 'inplace')
+        if not isinstance(keys, list):
+            keys = [keys]
+
+        if inplace:
+            obj = self
+        else:
+            obj = self.copy()
+
+        arrays = []
+        names = []
+        if append:
+            names = [x for x in self.index.names]
+            if isinstance(self.index, MultiIndex):
+                for i in range(self.index.nlevels):
+                    arrays.append(self.index._get_level_values(i))
+            else:
+                arrays.append(self.index)
+
+        to_remove = []
+        for col in keys:
+            if isinstance(col, MultiIndex):
+                # append all but the last column so we don't have to modify
+                # the end of this loop
+                for n in range(col.nlevels - 1):
+                    arrays.append(col._get_level_values(n))
+
+                level = col._get_level_values(col.nlevels - 1)
+                names.extend(col.names)
+            elif isinstance(col, Series):
+                level = col._values
+                names.append(col.name)
+            elif isinstance(col, Index):
+                level = col
+                names.append(col.name)
+            elif isinstance(col, (list, np.ndarray)):
+                level = col
+                names.append(None)
+            elif isinstance(obj, Series):
+                # col may not be a column label for Series case
+                raise ValueError('asdf')
+            else:
+                level = obj[col]._values
+                names.append(col)
+                if drop:
+                    to_remove.append(col)
+            arrays.append(level)
+
+        index = ensure_index_from_sequences(arrays, names)
+
+        if verify_integrity and not index.is_unique:
+            duplicates = index[index.duplicated()].unique()
+            raise ValueError('Index has duplicate keys: {dup}'.format(
+                dup=duplicates))
+
+        for c in to_remove:
+            del obj[c]
+
+        # clear up memory usage
+        index._cleanup()
+
+        obj.index = index
+
+        if not inplace:
+            return obj
+
     @Appender(_shared_docs['transpose'] % _shared_doc_kwargs)
     def transpose(self, *args, **kwargs):
 

pandas/core/series.py

@@ -1094,6 +1094,36 @@ def _set_value(self, label, value, takeable=False):
         return self
     _set_value.__doc__ = set_value.__doc__
 
+    def set_index(self, arrays, append=False, inplace=False,
+                  verify_integrity=False):
+        """
+        Set the Series index (row labels) using one or more columns.
+        By default yields a new object.
+
+        Parameters
+        ----------
+        arrays : array or list of arrays
+        append : boolean, default False
+            Whether to append columns to existing index
+        inplace : boolean, default False
+            Modify the Series in place (do not create a new object)
+        verify_integrity : boolean, default False
+            Check the new index for duplicates. Otherwise defer the check until
+            necessary. Setting to False will improve the performance of this
+            method
+
+        Returns
+        -------
+        reindexed : Series if inplace is False, else None
+
+        Examples
+        --------
+        >>> ...
+        """
+        return super(Series, self).set_index(keys=arrays, drop=False,
+                                             append=append, inplace=inplace,
+                                             verify_integrity=verify_integrity)
+
     def reset_index(self, level=None, drop=False, name=None, inplace=False):
         """
         Generate a new DataFrame or Series with the index reset.

pandas/tests/frame/test_indexing.py

@@ -2320,6 +2320,9 @@ def test_at_time_between_time_datetimeindex(self):
         result.loc[bkey] = df.iloc[binds]
         assert_frame_equal(result, df)
 
+    def test_set_index(self):
+        pytest.skip('TODO')
+
     def test_xs(self):
         idx = self.frame.index[5]
         xs = self.frame.xs(idx)

pandas/tests/series/test_alter_axes.py

@@ -20,19 +20,11 @@
 
 class TestSeriesAlterAxes(TestData):
 
-    def test_setindex(self):
-        # wrong type
-        series = self.series.copy()
-        pytest.raises(TypeError, setattr, series, 'index', None)
-
-        # wrong length
-        series = self.series.copy()
-        pytest.raises(Exception, setattr, series, 'index',
-                      np.arange(len(series) - 1))
+    def test_set_index(self):
+        # TODO!
 
-        # works
         series = self.series.copy()
-        series.index = np.arange(len(series))
+        series = series.set_index(np.arange(len(series)))
         assert isinstance(series.index, Index)
 
     def test_rename(self):
@@ -116,7 +108,7 @@ def test_set_index_makes_timeseries(self):
         idx = tm.makeDateIndex(10)
 
         s = Series(lrange(10))
-        s.index = idx
+        s = s.set_index(idx)
         assert s.index.is_all_dates
 
     def test_reset_index(self):