Skip to content

Commit

Permalink
ENH: ExtensionArray.repeat (pandas-dev#24349)
Browse files Browse the repository at this point in the history
  • Loading branch information
TomAugspurger authored and Pingviinituutti committed Feb 28, 2019
1 parent 8510256 commit f8b0466
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 1 deletion.
3 changes: 2 additions & 1 deletion doc/source/whatsnew/v0.24.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -987,7 +987,8 @@ update the ``ExtensionDtype._metadata`` tuple to match the signature of your

**Other changes**

- ``ExtensionArray`` has gained the abstract methods ``.dropna()`` (:issue:`21185`)
- :meth:`~pandas.api.types.ExtensionArray.dropna` has been added (:issue:`21185`)
- :meth:`~pandas.api.types.ExtensionArray.repeat` has been added (:issue:`24349`)
- ``ExtensionDtype`` has gained the ability to instantiate from string dtypes, e.g. ``decimal`` would instantiate a registered ``DecimalDtype``; furthermore
the ``ExtensionDtype`` has gained the method ``construct_array_type`` (:issue:`21185`)
- An ``ExtensionArray`` with a boolean dtype now works correctly as a boolean indexer. :meth:`pandas.api.types.is_bool_dtype` now properly considers them boolean (:issue:`22326`)
Expand Down
29 changes: 29 additions & 0 deletions pandas/core/arrays/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -580,6 +580,35 @@ def factorize(self, na_sentinel=-1):
uniques = self._from_factorized(uniques, self)
return labels, uniques

def repeat(self, repeats, axis=None):
"""
Repeat elements of an array.
.. versionadded:: 0.24.0
Parameters
----------
repeats : int
This should be a non-negative integer. Repeating 0 times
will return an empty array.
Returns
-------
repeated_array : ExtensionArray
Same type as the input, with elements repeated `repeats` times.
See Also
--------
numpy.repeat : Similar method for :class:`numpy.ndarray`.
ExtensionArray.take : Take arbitrary positions.
"""
if axis is not None:
raise ValueError("'axis' must be None.")
if repeats < 0:
raise ValueError("negative repeats are not allowed.")
ind = np.arange(len(self)).repeat(repeats)
return self.take(ind)

# ------------------------------------------------------------------------
# Indexing methods
# ------------------------------------------------------------------------
Expand Down
31 changes: 31 additions & 0 deletions pandas/tests/extension/base/methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,3 +264,34 @@ def test_where_series(self, data, na_value, as_frame):
if as_frame:
expected = expected.to_frame(name='a')
self.assert_equal(result, expected)

@pytest.mark.parametrize("as_series", [True, False])
@pytest.mark.parametrize("repeats", [0, 1, 2])
def test_repeat(self, data, repeats, as_series):
a, b, c = data[:3]
arr = type(data)._from_sequence([a, b, c], dtype=data.dtype)

if as_series:
arr = pd.Series(arr)

result = arr.repeat(repeats)

if repeats == 0:
expected = []
elif repeats == 1:
expected = [a, b, c]
else:
expected = [a, a, b, b, c, c]
expected = type(data)._from_sequence(expected, dtype=data.dtype)
if as_series:
index = pd.Series(np.arange(len(arr))).repeat(repeats).index
expected = pd.Series(expected, index=index)
self.assert_equal(result, expected)

def test_repeat_raises(self, data):
with pytest.raises(ValueError, match="'axis'"):
data.repeat(2, axis=1)

with pytest.raises(ValueError,
match="negative"):
data.repeat(-1)

0 comments on commit f8b0466

Please sign in to comment.