Skip to content

Commit

Permalink
ENH: Make RangeIndex.append() return RangeIndex when possible
Browse files Browse the repository at this point in the history
  • Loading branch information
toobaz committed May 3, 2017
1 parent f154966 commit 3a1f6dd
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 0 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.20.1.txt
Expand Up @@ -18,6 +18,7 @@ Highlights include:

Enhancements
~~~~~~~~~~~~
- ``RangeIndex.append`` now returns a ``RangeIndex`` object when possible (:issue:`16212`)



Expand Down
57 changes: 57 additions & 0 deletions pandas/core/indexes/range.py
Expand Up @@ -443,6 +443,63 @@ def join(self, other, how='left', level=None, return_indexers=False,
return super(RangeIndex, self).join(other, how, level, return_indexers,
sort)

def append(self, other):
"""
Append a collection of Index options together
Parameters
----------
other : Index or list/tuple of indices
Returns
-------
appended : RangeIndex if all indexes are consecutive RangeIndexes,
otherwise Int64Index or Index
"""

to_concat = [self]

if isinstance(other, (list, tuple)):
to_concat = to_concat + list(other)
else:
to_concat.append(other)

if not all([isinstance(i, RangeIndex) for i in to_concat]):
return super(RangeIndex, self).append(other)

start = step = next = None

for obj in to_concat:
if not len(obj):
continue

if start is None:
# This is set by the first non-empty index
start = obj._start
if step is None and len(obj) > 1:
step = obj._step
elif step is None:
# First non-empty index had only one element
if obj._start == start:
return super(RangeIndex, self).append(other)
step = obj._start - start

non_consecutive = ((step != obj._step and len(obj) > 1) or
(next is not None and obj._start != next))
if non_consecutive:
return super(RangeIndex, self).append(other)

if step is not None:
next = obj[-1] + step

if start is None:
start = obj._start
step = obj._step
stop = obj._stop if next is None else next
names = set([obj.name for obj in to_concat])
name = None if len(names) > 1 else self.name
return RangeIndex(start, stop, step, name=name)

def __len__(self):
"""
return the length of the RangeIndex
Expand Down
38 changes: 38 additions & 0 deletions pandas/tests/indexes/test_range.py
Expand Up @@ -941,3 +941,41 @@ def test_where_array_like(self):
for klass in klasses:
result = i.where(klass(cond))
tm.assert_index_equal(result, expected)

def test_append(self):
RI = RangeIndex
I64 = Int64Index
F64 = Float64Index
OI = Index
cases = [([RI(1, 12, 5)], RI(1, 12, 5)),
([RI(0, 6, 4)], RI(0, 6, 4)),
([RI(1, 3), RI(3, 7)], RI(1, 7)),
([RI(1, 5, 2), RI(5, 6)], RI(1, 6, 2)),
([RI(1, 3, 2), RI(4, 7, 3)], RI(1, 7, 3)),
([RI(-4, 3, 2), RI(4, 7, 2)], RI(-4, 7, 2)),
([RI(-4, -8), RI(-8, -12)], RI(-8, -12)),
([RI(-4, -8), RI(3, -4)], RI(3, -8)),
([RI(-4, -8), RI(3, 5)], RI(3, 5)),
([RI(-4, -2), RI(3, 5)], I64([-4, -3, 3, 4])),
([RI(-2,), RI(3, 5)], RI(3, 5)),
([RI(2,), RI(2)], I64([0, 1, 0, 1])),
([RI(2,), RI(2, 5), RI(5, 8, 4)], RI(0, 6)),
([RI(2,), RI(3, 5), RI(5, 8, 4)], I64([0, 1, 3, 4, 5])),
([RI(-2, 2), RI(2, 5), RI(5, 8, 4)], RI(-2, 6)),
([RI(3,), pd.Int64Index([-1, 3, 15])],
I64([0, 1, 2, -1, 3, 15])),
([RI(3,), pd.Float64Index([-1, 3.1, 15.0])],
F64([0, 1, 2, -1, 3.1, 15.0])),
([RI(3,), pd.Index(['a', None, 14])],
OI([0, 1, 2, 'a', None, 14])),
([RI(3, 1), pd.Index(['a', None, 14])], OI(['a', None, 14]))
]

for indices, expected in cases:
result = indices[0].append(indices[1:])
tm.assert_index_equal(result, expected, exact=True)

if len(indices) == 2:
# Append single item rather than list
result2 = indices[0].append(indices[1])
tm.assert_index_equal(result2, expected, exact=True)

0 comments on commit 3a1f6dd

Please sign in to comment.