Skip to content

Commit 52538fa

Browse files
committed
BUG: divmod return type
1 parent a277e4a commit 52538fa

File tree

4 files changed

+50
-7
lines changed

4 files changed

+50
-7
lines changed

doc/source/extending.rst

+12-3
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,18 @@ your ``MyExtensionArray`` class, as follows:
160160
MyExtensionArray._add_arithmetic_ops()
161161
MyExtensionArray._add_comparison_ops()
162162
163-
Note that since ``pandas`` automatically calls the underlying operator on each
164-
element one-by-one, this might not be as performant as implementing your own
165-
version of the associated operators directly on the ``ExtensionArray``.
163+
164+
.. note::
165+
166+
Since ``pandas`` automatically calls the underlying operator on each
167+
element one-by-one, this might not be as performant as implementing your own
168+
version of the associated operators directly on the ``ExtensionArray``.
169+
170+
This implementation will try to reconstruct a new ``ExtensionArray`` with the
171+
result of the element-wise operation. Whether or not that succeeds depends on
172+
whether the operation returns a result that's valid for the ``ExtensionArray``.
173+
If an ``ExtensionArray`` cannot be reconstructed, a list containing the scalars
174+
returned instead.
166175

167176
.. _extending.extension.testing:
168177

pandas/core/arrays/base.py

+12-4
Original file line numberDiff line numberDiff line change
@@ -775,10 +775,18 @@ def convert_values(param):
775775
res = [op(a, b) for (a, b) in zip(lvalues, rvalues)]
776776

777777
if coerce_to_dtype:
778-
try:
779-
res = self._from_sequence(res)
780-
except TypeError:
781-
pass
778+
if op.__name__ in {'divmod', 'rdivmod'}:
779+
try:
780+
a, b = zip(*res)
781+
res = (self._from_sequence(a),
782+
self._from_sequence(b))
783+
except TypeError:
784+
pass
785+
else:
786+
try:
787+
res = self._from_sequence(res)
788+
except TypeError:
789+
pass
782790

783791
return res
784792

pandas/tests/extension/decimal/array.py

+4
Original file line numberDiff line numberDiff line change
@@ -138,5 +138,9 @@ def _concat_same_type(cls, to_concat):
138138
return cls(np.concatenate([x._data for x in to_concat]))
139139

140140

141+
def to_decimal(values, context=None):
142+
return DecimalArray([decimal.Decimal(x) for x in values], context=context)
143+
144+
141145
DecimalArray._add_arithmetic_ops()
142146
DecimalArray._add_comparison_ops()

pandas/tests/extension/test_ops.py

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import pytest
2+
3+
from pandas.tests.extension.decimal.array import to_decimal
4+
import pandas.util.testing as tm
5+
6+
7+
@pytest.mark.parametrize("reverse, expected_div, expected_mod", [
8+
(False, [0, 1, 1, 2], [1, 0, 1, 0]),
9+
(True, [2, 1, 0, 0], [0, 0, 2, 2]),
10+
])
11+
def test_divmod(reverse, expected_div, expected_mod):
12+
# https://github.com/pandas-dev/pandas/issues/22930
13+
arr = to_decimal([1, 2, 3, 4])
14+
if reverse:
15+
div, mod = divmod(2, arr)
16+
else:
17+
div, mod = divmod(arr, 2)
18+
expected_div = to_decimal(expected_div)
19+
expected_mod = to_decimal(expected_mod)
20+
21+
tm.assert_extension_array_equal(div, expected_div)
22+
tm.assert_extension_array_equal(mod, expected_mod)

0 commit comments

Comments
 (0)