1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
from __future__ import annotations
 
import decimal
import operator
 
import numpy as np
import pytest
 
from pandas.compat.numpy import np_version_gt2
 
import pandas as pd
import pandas._testing as tm
from pandas.tests.extension import base
from pandas.tests.extension.decimal.array import (
    DecimalArray,
    DecimalDtype,
    make_data,
    to_decimal,
)
 
 
@pytest.fixture
def dtype():
    return DecimalDtype()
 
 
@pytest.fixture
def data():
    return DecimalArray(make_data())
 
 
@pytest.fixture
def data_for_twos():
    return DecimalArray([decimal.Decimal(2) for _ in range(100)])
 
 
@pytest.fixture
def data_missing():
    return DecimalArray([decimal.Decimal("NaN"), decimal.Decimal(1)])
 
 
@pytest.fixture
def data_for_sorting():
    return DecimalArray(
        [decimal.Decimal("1"), decimal.Decimal("2"), decimal.Decimal("0")]
    )
 
 
@pytest.fixture
def data_missing_for_sorting():
    return DecimalArray(
        [decimal.Decimal("1"), decimal.Decimal("NaN"), decimal.Decimal("0")]
    )
 
 
@pytest.fixture
def na_cmp():
    return lambda x, y: x.is_nan() and y.is_nan()
 
 
@pytest.fixture
def data_for_grouping():
    b = decimal.Decimal("1.0")
    a = decimal.Decimal("0.0")
    c = decimal.Decimal("2.0")
    na = decimal.Decimal("NaN")
    return DecimalArray([b, b, na, na, a, a, b, c])
 
 
class TestDecimalArray(base.ExtensionTests):
    def _get_expected_exception(
        self, op_name: str, obj, other
    ) -> type[Exception] | tuple[type[Exception], ...] | None:
        return None
 
    def _supports_reduction(self, ser: pd.Series, op_name: str) -> bool:
        return True
 
    def check_reduce(self, ser: pd.Series, op_name: str, skipna: bool):
        if op_name == "count":
            return super().check_reduce(ser, op_name, skipna)
        else:
            result = getattr(ser, op_name)(skipna=skipna)
            expected = getattr(np.asarray(ser), op_name)()
            tm.assert_almost_equal(result, expected)
 
    def test_reduce_series_numeric(self, data, all_numeric_reductions, skipna, request):
        if all_numeric_reductions in ["kurt", "skew", "sem", "median"]:
            mark = pytest.mark.xfail(raises=NotImplementedError)
            request.applymarker(mark)
        super().test_reduce_series_numeric(data, all_numeric_reductions, skipna)
 
    def test_reduce_frame(self, data, all_numeric_reductions, skipna, request):
        op_name = all_numeric_reductions
        if op_name in ["skew", "median"]:
            mark = pytest.mark.xfail(raises=NotImplementedError)
            request.applymarker(mark)
 
        return super().test_reduce_frame(data, all_numeric_reductions, skipna)
 
    def test_compare_scalar(self, data, comparison_op):
        ser = pd.Series(data)
        self._compare_other(ser, data, comparison_op, 0.5)
 
    def test_compare_array(self, data, comparison_op):
        ser = pd.Series(data)
 
        alter = np.random.default_rng(2).choice([-1, 0, 1], len(data))
        # Randomly double, halve or keep same value
        other = pd.Series(data) * [decimal.Decimal(pow(2.0, i)) for i in alter]
        self._compare_other(ser, data, comparison_op, other)
 
    def test_arith_series_with_array(self, data, all_arithmetic_operators):
        op_name = all_arithmetic_operators
        ser = pd.Series(data)
 
        context = decimal.getcontext()
        divbyzerotrap = context.traps[decimal.DivisionByZero]
        invalidoptrap = context.traps[decimal.InvalidOperation]
        context.traps[decimal.DivisionByZero] = 0
        context.traps[decimal.InvalidOperation] = 0
 
        # Decimal supports ops with int, but not float
        other = pd.Series([int(d * 100) for d in data])
        self.check_opname(ser, op_name, other)
 
        if "mod" not in op_name:
            self.check_opname(ser, op_name, ser * 2)
 
        self.check_opname(ser, op_name, 0)
        self.check_opname(ser, op_name, 5)
        context.traps[decimal.DivisionByZero] = divbyzerotrap
        context.traps[decimal.InvalidOperation] = invalidoptrap
 
    def test_fillna_frame(self, data_missing):
        msg = "ExtensionArray.fillna added a 'copy' keyword"
        with tm.assert_produces_warning(
            DeprecationWarning, match=msg, check_stacklevel=False
        ):
            super().test_fillna_frame(data_missing)
 
    def test_fillna_limit_pad(self, data_missing):
        msg = "ExtensionArray.fillna 'method' keyword is deprecated"
        with tm.assert_produces_warning(
            DeprecationWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            super().test_fillna_limit_pad(data_missing)
 
        msg = "The 'method' keyword in DecimalArray.fillna is deprecated"
        with tm.assert_produces_warning(
            FutureWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            super().test_fillna_limit_pad(data_missing)
 
    @pytest.mark.parametrize(
        "limit_area, input_ilocs, expected_ilocs",
        [
            ("outside", [1, 0, 0, 0, 1], [1, 0, 0, 0, 1]),
            ("outside", [1, 0, 1, 0, 1], [1, 0, 1, 0, 1]),
            ("outside", [0, 1, 1, 1, 0], [0, 1, 1, 1, 1]),
            ("outside", [0, 1, 0, 1, 0], [0, 1, 0, 1, 1]),
            ("inside", [1, 0, 0, 0, 1], [1, 1, 1, 1, 1]),
            ("inside", [1, 0, 1, 0, 1], [1, 1, 1, 1, 1]),
            ("inside", [0, 1, 1, 1, 0], [0, 1, 1, 1, 0]),
            ("inside", [0, 1, 0, 1, 0], [0, 1, 1, 1, 0]),
        ],
    )
    def test_ffill_limit_area(
        self, data_missing, limit_area, input_ilocs, expected_ilocs
    ):
        # GH#56616
        msg = "ExtensionArray.fillna 'method' keyword is deprecated"
        with tm.assert_produces_warning(
            DeprecationWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            msg = "DecimalArray does not implement limit_area"
            with pytest.raises(NotImplementedError, match=msg):
                super().test_ffill_limit_area(
                    data_missing, limit_area, input_ilocs, expected_ilocs
                )
 
    def test_fillna_limit_backfill(self, data_missing):
        msg = "Series.fillna with 'method' is deprecated"
        with tm.assert_produces_warning(
            FutureWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            super().test_fillna_limit_backfill(data_missing)
 
        msg = "ExtensionArray.fillna 'method' keyword is deprecated"
        with tm.assert_produces_warning(
            DeprecationWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            super().test_fillna_limit_backfill(data_missing)
 
        msg = "The 'method' keyword in DecimalArray.fillna is deprecated"
        with tm.assert_produces_warning(
            FutureWarning,
            match=msg,
            check_stacklevel=False,
            raise_on_extra_warnings=False,
        ):
            super().test_fillna_limit_backfill(data_missing)
 
    def test_fillna_no_op_returns_copy(self, data):
        msg = "|".join(
            [
                "ExtensionArray.fillna 'method' keyword is deprecated",
                "The 'method' keyword in DecimalArray.fillna is deprecated",
            ]
        )
        with tm.assert_produces_warning(
            (FutureWarning, DeprecationWarning), match=msg, check_stacklevel=False
        ):
            super().test_fillna_no_op_returns_copy(data)
 
    def test_fillna_series(self, data_missing):
        msg = "ExtensionArray.fillna added a 'copy' keyword"
        with tm.assert_produces_warning(
            DeprecationWarning, match=msg, check_stacklevel=False
        ):
            super().test_fillna_series(data_missing)
 
    def test_fillna_series_method(self, data_missing, fillna_method):
        msg = "|".join(
            [
                "ExtensionArray.fillna 'method' keyword is deprecated",
                "The 'method' keyword in DecimalArray.fillna is deprecated",
            ]
        )
        with tm.assert_produces_warning(
            (FutureWarning, DeprecationWarning), match=msg, check_stacklevel=False
        ):
            super().test_fillna_series_method(data_missing, fillna_method)
 
    def test_fillna_copy_frame(self, data_missing, using_copy_on_write):
        warn = DeprecationWarning if not using_copy_on_write else None
        msg = "ExtensionArray.fillna added a 'copy' keyword"
        with tm.assert_produces_warning(warn, match=msg, check_stacklevel=False):
            super().test_fillna_copy_frame(data_missing)
 
    def test_fillna_copy_series(self, data_missing, using_copy_on_write):
        warn = DeprecationWarning if not using_copy_on_write else None
        msg = "ExtensionArray.fillna added a 'copy' keyword"
        with tm.assert_produces_warning(warn, match=msg, check_stacklevel=False):
            super().test_fillna_copy_series(data_missing)
 
    @pytest.mark.parametrize("dropna", [True, False])
    def test_value_counts(self, all_data, dropna, request):
        all_data = all_data[:10]
        if dropna:
            other = np.array(all_data[~all_data.isna()])
        else:
            other = all_data
 
        vcs = pd.Series(all_data).value_counts(dropna=dropna)
        vcs_ex = pd.Series(other).value_counts(dropna=dropna)
 
        with decimal.localcontext() as ctx:
            # avoid raising when comparing Decimal("NAN") < Decimal(2)
            ctx.traps[decimal.InvalidOperation] = False
 
            result = vcs.sort_index()
            expected = vcs_ex.sort_index()
 
        tm.assert_series_equal(result, expected)
 
    def test_series_repr(self, data):
        # Overriding this base test to explicitly test that
        # the custom _formatter is used
        ser = pd.Series(data)
        assert data.dtype.name in repr(ser)
        assert "Decimal: " in repr(ser)
 
    @pytest.mark.xfail(reason="Inconsistent array-vs-scalar behavior")
    @pytest.mark.parametrize("ufunc", [np.positive, np.negative, np.abs])
    def test_unary_ufunc_dunder_equivalence(self, data, ufunc):
        super().test_unary_ufunc_dunder_equivalence(data, ufunc)
 
    def test_array_interface_copy(self, data):
        result_copy1 = np.array(data, copy=True)
        result_copy2 = np.array(data, copy=True)
        assert not np.may_share_memory(result_copy1, result_copy2)
        if not np_version_gt2:
            # copy=False semantics are only supported in NumPy>=2.
            return
 
        try:
            result_nocopy1 = np.array(data, copy=False)
        except ValueError:
            # An error is always acceptable for `copy=False`
            return
 
        result_nocopy2 = np.array(data, copy=False)
        # If copy=False was given and did not raise, these must share the same data
        assert np.may_share_memory(result_nocopy1, result_nocopy2)
 
 
def test_take_na_value_other_decimal():
    arr = DecimalArray([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
    result = arr.take([0, -1], allow_fill=True, fill_value=decimal.Decimal("-1.0"))
    expected = DecimalArray([decimal.Decimal("1.0"), decimal.Decimal("-1.0")])
    tm.assert_extension_array_equal(result, expected)
 
 
def test_series_constructor_coerce_data_to_extension_dtype():
    dtype = DecimalDtype()
    ser = pd.Series([0, 1, 2], dtype=dtype)
 
    arr = DecimalArray(
        [decimal.Decimal(0), decimal.Decimal(1), decimal.Decimal(2)],
        dtype=dtype,
    )
    exp = pd.Series(arr)
    tm.assert_series_equal(ser, exp)
 
 
def test_series_constructor_with_dtype():
    arr = DecimalArray([decimal.Decimal("10.0")])
    result = pd.Series(arr, dtype=DecimalDtype())
    expected = pd.Series(arr)
    tm.assert_series_equal(result, expected)
 
    result = pd.Series(arr, dtype="int64")
    expected = pd.Series([10])
    tm.assert_series_equal(result, expected)
 
 
def test_dataframe_constructor_with_dtype():
    arr = DecimalArray([decimal.Decimal("10.0")])
 
    result = pd.DataFrame({"A": arr}, dtype=DecimalDtype())
    expected = pd.DataFrame({"A": arr})
    tm.assert_frame_equal(result, expected)
 
    arr = DecimalArray([decimal.Decimal("10.0")])
    result = pd.DataFrame({"A": arr}, dtype="int64")
    expected = pd.DataFrame({"A": [10]})
    tm.assert_frame_equal(result, expected)
 
 
@pytest.mark.parametrize("frame", [True, False])
def test_astype_dispatches(frame):
    # This is a dtype-specific test that ensures Series[decimal].astype
    # gets all the way through to ExtensionArray.astype
    # Designing a reliable smoke test that works for arbitrary data types
    # is difficult.
    data = pd.Series(DecimalArray([decimal.Decimal(2)]), name="a")
    ctx = decimal.Context()
    ctx.prec = 5
 
    if frame:
        data = data.to_frame()
 
    result = data.astype(DecimalDtype(ctx))
 
    if frame:
        result = result["a"]
 
    assert result.dtype.context.prec == ctx.prec
 
 
class DecimalArrayWithoutFromSequence(DecimalArray):
    """Helper class for testing error handling in _from_sequence."""
 
    @classmethod
    def _from_sequence(cls, scalars, *, dtype=None, copy=False):
        raise KeyError("For the test")
 
 
class DecimalArrayWithoutCoercion(DecimalArrayWithoutFromSequence):
    @classmethod
    def _create_arithmetic_method(cls, op):
        return cls._create_method(op, coerce_to_dtype=False)
 
 
DecimalArrayWithoutCoercion._add_arithmetic_ops()
 
 
def test_combine_from_sequence_raises(monkeypatch):
    # https://github.com/pandas-dev/pandas/issues/22850
    cls = DecimalArrayWithoutFromSequence
 
    @classmethod
    def construct_array_type(cls):
        return DecimalArrayWithoutFromSequence
 
    monkeypatch.setattr(DecimalDtype, "construct_array_type", construct_array_type)
 
    arr = cls([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
    ser = pd.Series(arr)
    result = ser.combine(ser, operator.add)
 
    # note: object dtype
    expected = pd.Series(
        [decimal.Decimal("2.0"), decimal.Decimal("4.0")], dtype="object"
    )
    tm.assert_series_equal(result, expected)
 
 
@pytest.mark.parametrize(
    "class_", [DecimalArrayWithoutFromSequence, DecimalArrayWithoutCoercion]
)
def test_scalar_ops_from_sequence_raises(class_):
    # op(EA, EA) should return an EA, or an ndarray if it's not possible
    # to return an EA with the return values.
    arr = class_([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
    result = arr + arr
    expected = np.array(
        [decimal.Decimal("2.0"), decimal.Decimal("4.0")], dtype="object"
    )
    tm.assert_numpy_array_equal(result, expected)
 
 
@pytest.mark.parametrize(
    "reverse, expected_div, expected_mod",
    [(False, [0, 1, 1, 2], [1, 0, 1, 0]), (True, [2, 1, 0, 0], [0, 0, 2, 2])],
)
def test_divmod_array(reverse, expected_div, expected_mod):
    # https://github.com/pandas-dev/pandas/issues/22930
    arr = to_decimal([1, 2, 3, 4])
    if reverse:
        div, mod = divmod(2, arr)
    else:
        div, mod = divmod(arr, 2)
    expected_div = to_decimal(expected_div)
    expected_mod = to_decimal(expected_mod)
 
    tm.assert_extension_array_equal(div, expected_div)
    tm.assert_extension_array_equal(mod, expected_mod)
 
 
def test_ufunc_fallback(data):
    a = data[:5]
    s = pd.Series(a, index=range(3, 8))
    result = np.abs(s)
    expected = pd.Series(np.abs(a), index=range(3, 8))
    tm.assert_series_equal(result, expected)
 
 
def test_array_ufunc():
    a = to_decimal([1, 2, 3])
    result = np.exp(a)
    expected = to_decimal(np.exp(a._data))
    tm.assert_extension_array_equal(result, expected)
 
 
def test_array_ufunc_series():
    a = to_decimal([1, 2, 3])
    s = pd.Series(a)
    result = np.exp(s)
    expected = pd.Series(to_decimal(np.exp(a._data)))
    tm.assert_series_equal(result, expected)
 
 
def test_array_ufunc_series_scalar_other():
    # check _HANDLED_TYPES
    a = to_decimal([1, 2, 3])
    s = pd.Series(a)
    result = np.add(s, decimal.Decimal(1))
    expected = pd.Series(np.add(a, decimal.Decimal(1)))
    tm.assert_series_equal(result, expected)
 
 
def test_array_ufunc_series_defer():
    a = to_decimal([1, 2, 3])
    s = pd.Series(a)
 
    expected = pd.Series(to_decimal([2, 4, 6]))
    r1 = np.add(s, a)
    r2 = np.add(a, s)
 
    tm.assert_series_equal(r1, expected)
    tm.assert_series_equal(r2, expected)
 
 
def test_groupby_agg():
    # Ensure that the result of agg is inferred to be decimal dtype
    # https://github.com/pandas-dev/pandas/issues/29141
 
    data = make_data()[:5]
    df = pd.DataFrame(
        {"id1": [0, 0, 0, 1, 1], "id2": [0, 1, 0, 1, 1], "decimals": DecimalArray(data)}
    )
 
    # single key, selected column
    expected = pd.Series(to_decimal([data[0], data[3]]))
    result = df.groupby("id1")["decimals"].agg(lambda x: x.iloc[0])
    tm.assert_series_equal(result, expected, check_names=False)
    result = df["decimals"].groupby(df["id1"]).agg(lambda x: x.iloc[0])
    tm.assert_series_equal(result, expected, check_names=False)
 
    # multiple keys, selected column
    expected = pd.Series(
        to_decimal([data[0], data[1], data[3]]),
        index=pd.MultiIndex.from_tuples([(0, 0), (0, 1), (1, 1)]),
    )
    result = df.groupby(["id1", "id2"])["decimals"].agg(lambda x: x.iloc[0])
    tm.assert_series_equal(result, expected, check_names=False)
    result = df["decimals"].groupby([df["id1"], df["id2"]]).agg(lambda x: x.iloc[0])
    tm.assert_series_equal(result, expected, check_names=False)
 
    # multiple columns
    expected = pd.DataFrame({"id2": [0, 1], "decimals": to_decimal([data[0], data[3]])})
    result = df.groupby("id1").agg(lambda x: x.iloc[0])
    tm.assert_frame_equal(result, expected, check_names=False)
 
 
def test_groupby_agg_ea_method(monkeypatch):
    # Ensure that the result of agg is inferred to be decimal dtype
    # https://github.com/pandas-dev/pandas/issues/29141
 
    def DecimalArray__my_sum(self):
        return np.sum(np.array(self))
 
    monkeypatch.setattr(DecimalArray, "my_sum", DecimalArray__my_sum, raising=False)
 
    data = make_data()[:5]
    df = pd.DataFrame({"id": [0, 0, 0, 1, 1], "decimals": DecimalArray(data)})
    expected = pd.Series(to_decimal([data[0] + data[1] + data[2], data[3] + data[4]]))
 
    result = df.groupby("id")["decimals"].agg(lambda x: x.values.my_sum())
    tm.assert_series_equal(result, expected, check_names=False)
    s = pd.Series(DecimalArray(data))
    grouper = np.array([0, 0, 0, 1, 1], dtype=np.int64)
    result = s.groupby(grouper).agg(lambda x: x.values.my_sum())
    tm.assert_series_equal(result, expected, check_names=False)
 
 
def test_indexing_no_materialize(monkeypatch):
    # See https://github.com/pandas-dev/pandas/issues/29708
    # Ensure that indexing operations do not materialize (convert to a numpy
    # array) the ExtensionArray unnecessary
 
    def DecimalArray__array__(self, dtype=None):
        raise Exception("tried to convert a DecimalArray to a numpy array")
 
    monkeypatch.setattr(DecimalArray, "__array__", DecimalArray__array__, raising=False)
 
    data = make_data()
    s = pd.Series(DecimalArray(data))
    df = pd.DataFrame({"a": s, "b": range(len(s))})
 
    # ensure the following operations do not raise an error
    s[s > 0.5]
    df[s > 0.5]
    s.at[0]
    df.at[0, "a"]
 
 
def test_to_numpy_keyword():
    # test the extra keyword
    values = [decimal.Decimal("1.1111"), decimal.Decimal("2.2222")]
    expected = np.array(
        [decimal.Decimal("1.11"), decimal.Decimal("2.22")], dtype="object"
    )
    a = pd.array(values, dtype="decimal")
    result = a.to_numpy(decimals=2)
    tm.assert_numpy_array_equal(result, expected)
 
    result = pd.Series(a).to_numpy(decimals=2)
    tm.assert_numpy_array_equal(result, expected)
 
 
def test_array_copy_on_write(using_copy_on_write):
    df = pd.DataFrame({"a": [decimal.Decimal(2), decimal.Decimal(3)]}, dtype="object")
    df2 = df.astype(DecimalDtype())
    df.iloc[0, 0] = 0
    if using_copy_on_write:
        expected = pd.DataFrame(
            {"a": [decimal.Decimal(2), decimal.Decimal(3)]}, dtype=DecimalDtype()
        )
        tm.assert_equal(df2.values, expected.values)