ENH: RangeIndex, test suite passes now

Author: wesm

pandas/core/frame.py

@@ -2432,7 +2432,7 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
                 for i in range(self.index.nlevels):
                     arrays.append(self.index.get_level_values(i))
             else:
-                arrays.append(np.asarray(self.index))
+                arrays.append(self.index.values)
 
         to_remove = []
         for col in keys:

pandas/core/index.py

@@ -17,7 +17,7 @@
 __all__ = ['Index']
 
 
-def _indexOp(opname):
+def _index_comp(opname):
     """
     Wrapper function for index comparison operations, to avoid
     code duplication.
@@ -501,12 +501,12 @@ def __add__(self, other):
         else:
             return Index(self.view(np.ndarray) + other)
 
-    __eq__ = _indexOp('__eq__')
-    __ne__ = _indexOp('__ne__')
-    __lt__ = _indexOp('__lt__')
-    __gt__ = _indexOp('__gt__')
-    __le__ = _indexOp('__le__')
-    __ge__ = _indexOp('__ge__')
+    __eq__ = _index_comp('__eq__')
+    __ne__ = _index_comp('__ne__')
+    __lt__ = _index_comp('__lt__')
+    __gt__ = _index_comp('__gt__')
+    __le__ = _index_comp('__le__')
+    __ge__ = _index_comp('__ge__')
 
     def __sub__(self, other):
         return self.diff(other)
@@ -1071,6 +1071,7 @@ def slice_locs(self, start=None, end=None):
         -----
         This function assumes that the data is sorted, so use at your own peril
         """
+        values = None
         if start is None:
             beg_slice = 0
         else:
@@ -1095,6 +1096,9 @@ def slice_locs(self, start=None, end=None):
                     end_slice += 1
             except KeyError:
                 if self.is_monotonic:
+                    if values is None:
+                        values = self.values
+
                     end_slice = self.searchsorted(end, side='right')
                 else:
                     raise
@@ -1236,6 +1240,28 @@ def _wrap_joined_index(self, joined, other):
         return Int64Index(joined, name=name)
 
 
+def _range_comp(opname):
+    """
+    Wrapper function for index comparison operations, to avoid
+    code duplication.
+    """
+    def wrapper(self, other):
+        func = getattr(self.values, opname)
+        result = func(other)
+        return result
+    return wrapper
+
+def _range_arith(opname):
+    """
+    Wrapper function for index comparison operations, to avoid
+    code duplication.
+    """
+    def wrapper(self, other):
+        func = getattr(self._to_dense(), opname)
+        return func(other)
+    return wrapper
+
+
 class RangeIndex(Int64Index):
     """
     Compactly represents regular integer range (instead of generating a full
@@ -1269,6 +1295,30 @@ def __array_finalize__(self, obj):
         self.step = getattr(obj, 'step', None)
         self.name = getattr(obj, 'name', None)
 
+    __eq__ = _range_comp('__eq__')
+    __ne__ = _range_comp('__ne__')
+    __lt__ = _range_comp('__lt__')
+    __gt__ = _range_comp('__gt__')
+    __le__ = _range_comp('__le__')
+    __ge__ = _range_comp('__ge__')
+
+    __add__ = _range_arith('__add__')
+    __sub__ = _range_arith('__sub__')
+    __mul__ = _range_arith('__mul__')
+    __truediv__ = _range_arith('__truediv__')
+    __floordiv__ = _range_arith('__floordiv__')
+    __radd__ = _range_arith('__radd__')
+    __rsub__ = _range_arith('__rsub__')
+    __rmul__ = _range_arith('__rmul__')
+    __rtruediv__ = _range_arith('__rtruediv__')
+    __rfloordiv__ = _range_arith('__rfloordiv__')
+
+    # Python 2 division operators
+    if not py3compat.PY3:
+        __div__ = _range_arith('__div__')
+        __rdiv__ = _range_arith('__rdiv__')
+        __idiv__ = __div__
+
     @property
     def values(self):
         return np.arange(self.start, self.stop, self.step)
@@ -1329,10 +1379,10 @@ def __getitem__(self, key):
             raise TypeError('Unhandled getitem type: %s' % type(key))
 
     def insert(self, loc, item):
-        return Int64Index(self.values, name=self.name).insert(loc, item)
+        return self._to_dense().insert(loc, item)
 
     def delete(self, loc):
-        return Int64Index(self.values, name=self.name).delete(loc)
+        return self._to_dense().delete(loc)
 
     def take(self, key):
         # XXX
@@ -1352,12 +1402,59 @@ def equals(self, other):
         """
         if not isinstance(other, RangeIndex):
             if isinstance(other, Index):
-                other = other.values
+                return other.equals(self._to_dense())
             return np.array_equal(self.values, other)
 
         return (self.start == other.start and self.stop == other.stop
                 and self.step == other.step)
 
+    def searchsorted(self, key, side='left'):
+        return self.values.searchsorted(key, side=side)
+
+    def _to_dense(self):
+        return Int64Index(self.values, name=self.name)
+
+    def union(self, other):
+        """
+        Specialized Index.union for RangeIndex
+        """
+        if not isinstance(other, RangeIndex):
+            return self._to_dense().union(other)
+
+        if self.step == other.step:
+            this = self
+            if self.start > other.start:
+                this, other = other, this
+
+            if this.stop > other.start:
+                return RangeIndex(this.start, max(other.stop, this.stop),
+                                  self.step)
+            else:
+                return self._to_dense().union(other._to_dense())
+        else:
+            return self._to_dense().union(other._to_dense())
+
+    def intersection(self, other):
+        """
+        Specialized Index.intersection for RangeIndex
+        """
+        if not isinstance(other, RangeIndex):
+            return self._to_dense().intersection(other)
+
+        if self.step == other.step:
+            this = self
+            if self.start > other.start:
+                this, other = other, this
+
+            if other.stop <= this.stop:
+                return RangeIndex(other.start, min(this.stop, other.stop),
+                                  self.step)
+            else:
+                return self._to_dense().intersection(other._to_dense())
+        else:
+            return self._to_dense().intersection(other._to_dense())
+
+
     @property
     def _constructor(self):
         return RangeIndex

pandas/core/series.py

@@ -15,10 +15,11 @@
 import numpy.ma as ma
 
 from pandas.core.common import (isnull, notnull, _is_bool_indexer,
-                                _default_index, _maybe_upcast,
+                                _maybe_upcast,
                                 _asarray_tuplesafe, is_integer_dtype)
 from pandas.core.index import (Index, MultiIndex, InvalidIndexError,
-                               _ensure_index, _handle_legacy_indexes)
+                               _ensure_index, _handle_legacy_indexes,
+                               RangeIndex)
 from pandas.core.indexing import _SeriesIndexer
 from pandas.tseries.index import DatetimeIndex
 from pandas.tseries.period import PeriodIndex, Period
@@ -340,7 +341,7 @@ def __new__(cls, data=None, index=None, dtype=None, name=None,
             return subarr
 
         if index is None:
-            index = _default_index(len(subarr))
+            index = RangeIndex(len(subarr))
 
         # Change the class of the array to be the subclass type.
         if index.is_all_dates:
@@ -363,7 +364,9 @@ def from_array(cls, arr, index=None, name=None, copy=False):
             arr = arr.copy()
 
         klass = Series
-        if index.is_all_dates:
+        if index is None:
+            index = RangeIndex(len(arr))
+        elif index.is_all_dates:
             if not isinstance(index, (DatetimeIndex, PeriodIndex)):
                 index = DatetimeIndex(index)
             klass = TimeSeries
@@ -1153,7 +1156,7 @@ def max(self, axis=None, out=None, skipna=True, level=None):
 
     @Substitution(name='standard deviation', shortname='stdev',
                   na_action=_doc_exclude_na, extras='')
-    @Appender(_stat_doc + 
+    @Appender(_stat_doc +
         """
         Normalized by N-1 (unbiased estimator).
         """)
@@ -1166,7 +1169,7 @@ def std(self, axis=None, dtype=None, out=None, ddof=1, skipna=True,
 
     @Substitution(name='variance', shortname='var',
                   na_action=_doc_exclude_na, extras='')
-    @Appender(_stat_doc + 
+    @Appender(_stat_doc +
         """
         Normalized by N-1 (unbiased estimator).
         """)

pandas/tests/test_frame.py

@@ -1841,14 +1841,14 @@ def test_constructor_ndarray(self):
 
         # automatic labeling
         frame = DataFrame(mat)
-        self.assert_(np.array_equal(frame.index, range(2)))
-        self.assert_(np.array_equal(frame.columns, range(3)))
+        self.assert_(frame.index.equals(range(2)))
+        self.assert_(frame.columns.equals(range(3)))
 
         frame = DataFrame(mat, index=[1, 2])
-        self.assert_(np.array_equal(frame.columns, range(3)))
+        self.assert_(frame.columns.equals(range(3)))
 
         frame = DataFrame(mat, columns=['A', 'B', 'C'])
-        self.assert_(np.array_equal(frame.index, range(2)))
+        self.assert_(frame.index.equals(range(2)))
 
         # 0-length axis
         frame = DataFrame(np.empty((0, 3)))
@@ -1899,14 +1899,14 @@ def test_constructor_maskedarray(self):
 
         # automatic labeling
         frame = DataFrame(mat)
-        self.assert_(np.array_equal(frame.index, range(2)))
-        self.assert_(np.array_equal(frame.columns, range(3)))
+        self.assert_(frame.index.equals(range(2)))
+        self.assert_(frame.columns.equals(range(3)))
 
         frame = DataFrame(mat, index=[1, 2])
-        self.assert_(np.array_equal(frame.columns, range(3)))
+        self.assert_(frame.columns.equals(range(3)))
 
         frame = DataFrame(mat, columns=['A', 'B', 'C'])
-        self.assert_(np.array_equal(frame.index, range(2)))
+        self.assert_(frame.index.equals(range(2)))
 
         # 0-length axis
         frame = DataFrame(ma.masked_all((0, 3)))
@@ -6677,8 +6677,7 @@ def test_reset_index(self):
         deleveled = self.frame.reset_index()
         self.assert_(np.array_equal(deleveled['index'],
                                     self.frame.index.values))
-        self.assert_(np.array_equal(deleveled.index,
-                                    np.arange(len(deleveled))))
+        self.assert_(deleveled.index.equals(np.arange(len(deleveled))))
 
         # preserve column names
         self.frame.columns.name = 'columns'

pandas/tests/test_index.py

@@ -962,6 +962,10 @@ def test_slice(self):
         index = RangeIndex(0, 1)
         self.assertTrue(index[:0].equals(RangeIndex(0, 0)))
 
+    def test_slice_locs(self):
+        index = RangeIndex(0, 2)
+        self.assertEqual(index.slice_locs(0, 3), (0, 2))
+
     def test_insert(self):
         index = RangeIndex(5, name='bar')
 
@@ -978,10 +982,23 @@ def test_delete(self):
         self.assertEqual(result.name, 'bar')
 
     def test_union(self):
-        pass
+        result = RangeIndex(4).union(RangeIndex(3))
+        self.assertTrue(result.equals(RangeIndex(4)))
+
+        result = RangeIndex(3).union(RangeIndex(4, 7))
+        exp = Index([0, 1, 2, 4, 5, 6])
+        self.assertTrue(result.equals(exp))
 
     def test_intersection(self):
-        pass
+        result = RangeIndex(4).intersection(RangeIndex(1, 3))
+        self.assertTrue(result.equals(RangeIndex(1, 3)))
+
+    def test_array_arithmetic(self):
+        index = RangeIndex(0, 11)
+
+        result = index // 5
+        expected = index._to_dense() // 5
+        self.assertTrue(result.equals(expected))
 
     def test_pickle(self):
         pass

pandas/tools/tests/test_merge.py

@@ -560,7 +560,7 @@ def test_join_sort(self):
 
         # smoke test
         joined = left.join(right, on='key', sort=False)
-        self.assert_(np.array_equal(joined.index, range(4)))
+        self.assert_(joined.index.equals(range(4)))
 
     def test_intelligently_handle_join_key(self):
         # #733, be a bit more 1337 about not returning unconsolidated DataFrame