Skip to content

Commit 3a1f6dd

Browse files
committed
ENH: Make RangeIndex.append() return RangeIndex when possible
closes #16212
1 parent f154966 commit 3a1f6dd

File tree

3 files changed

+96
-0
lines changed

3 files changed

+96
-0
lines changed

doc/source/whatsnew/v0.20.1.txt

+1
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ Highlights include:
1818

1919
Enhancements
2020
~~~~~~~~~~~~
21+
- ``RangeIndex.append`` now returns a ``RangeIndex`` object when possible (:issue:`16212`)
2122

2223

2324

pandas/core/indexes/range.py

+57
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,63 @@ def join(self, other, how='left', level=None, return_indexers=False,
443443
return super(RangeIndex, self).join(other, how, level, return_indexers,
444444
sort)
445445

446+
def append(self, other):
447+
"""
448+
Append a collection of Index options together
449+
450+
Parameters
451+
----------
452+
other : Index or list/tuple of indices
453+
454+
Returns
455+
-------
456+
appended : RangeIndex if all indexes are consecutive RangeIndexes,
457+
otherwise Int64Index or Index
458+
"""
459+
460+
to_concat = [self]
461+
462+
if isinstance(other, (list, tuple)):
463+
to_concat = to_concat + list(other)
464+
else:
465+
to_concat.append(other)
466+
467+
if not all([isinstance(i, RangeIndex) for i in to_concat]):
468+
return super(RangeIndex, self).append(other)
469+
470+
start = step = next = None
471+
472+
for obj in to_concat:
473+
if not len(obj):
474+
continue
475+
476+
if start is None:
477+
# This is set by the first non-empty index
478+
start = obj._start
479+
if step is None and len(obj) > 1:
480+
step = obj._step
481+
elif step is None:
482+
# First non-empty index had only one element
483+
if obj._start == start:
484+
return super(RangeIndex, self).append(other)
485+
step = obj._start - start
486+
487+
non_consecutive = ((step != obj._step and len(obj) > 1) or
488+
(next is not None and obj._start != next))
489+
if non_consecutive:
490+
return super(RangeIndex, self).append(other)
491+
492+
if step is not None:
493+
next = obj[-1] + step
494+
495+
if start is None:
496+
start = obj._start
497+
step = obj._step
498+
stop = obj._stop if next is None else next
499+
names = set([obj.name for obj in to_concat])
500+
name = None if len(names) > 1 else self.name
501+
return RangeIndex(start, stop, step, name=name)
502+
446503
def __len__(self):
447504
"""
448505
return the length of the RangeIndex

pandas/tests/indexes/test_range.py

+38
Original file line numberDiff line numberDiff line change
@@ -941,3 +941,41 @@ def test_where_array_like(self):
941941
for klass in klasses:
942942
result = i.where(klass(cond))
943943
tm.assert_index_equal(result, expected)
944+
945+
def test_append(self):
946+
RI = RangeIndex
947+
I64 = Int64Index
948+
F64 = Float64Index
949+
OI = Index
950+
cases = [([RI(1, 12, 5)], RI(1, 12, 5)),
951+
([RI(0, 6, 4)], RI(0, 6, 4)),
952+
([RI(1, 3), RI(3, 7)], RI(1, 7)),
953+
([RI(1, 5, 2), RI(5, 6)], RI(1, 6, 2)),
954+
([RI(1, 3, 2), RI(4, 7, 3)], RI(1, 7, 3)),
955+
([RI(-4, 3, 2), RI(4, 7, 2)], RI(-4, 7, 2)),
956+
([RI(-4, -8), RI(-8, -12)], RI(-8, -12)),
957+
([RI(-4, -8), RI(3, -4)], RI(3, -8)),
958+
([RI(-4, -8), RI(3, 5)], RI(3, 5)),
959+
([RI(-4, -2), RI(3, 5)], I64([-4, -3, 3, 4])),
960+
([RI(-2,), RI(3, 5)], RI(3, 5)),
961+
([RI(2,), RI(2)], I64([0, 1, 0, 1])),
962+
([RI(2,), RI(2, 5), RI(5, 8, 4)], RI(0, 6)),
963+
([RI(2,), RI(3, 5), RI(5, 8, 4)], I64([0, 1, 3, 4, 5])),
964+
([RI(-2, 2), RI(2, 5), RI(5, 8, 4)], RI(-2, 6)),
965+
([RI(3,), pd.Int64Index([-1, 3, 15])],
966+
I64([0, 1, 2, -1, 3, 15])),
967+
([RI(3,), pd.Float64Index([-1, 3.1, 15.0])],
968+
F64([0, 1, 2, -1, 3.1, 15.0])),
969+
([RI(3,), pd.Index(['a', None, 14])],
970+
OI([0, 1, 2, 'a', None, 14])),
971+
([RI(3, 1), pd.Index(['a', None, 14])], OI(['a', None, 14]))
972+
]
973+
974+
for indices, expected in cases:
975+
result = indices[0].append(indices[1:])
976+
tm.assert_index_equal(result, expected, exact=True)
977+
978+
if len(indices) == 2:
979+
# Append single item rather than list
980+
result2 = indices[0].append(indices[1])
981+
tm.assert_index_equal(result2, expected, exact=True)

0 commit comments

Comments
 (0)