Brettapps commited on
Commit
abd0e97
·
verified ·
1 Parent(s): 03de080

Upload folder using huggingface_hub (part 11)

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .venv/lib/python3.13/site-packages/pandas/tests/apply/common.py +7 -0
  2. .venv/lib/python3.13/site-packages/pandas/tests/apply/conftest.py +63 -0
  3. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_apply.py +1875 -0
  4. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_apply_relabeling.py +105 -0
  5. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_transform.py +264 -0
  6. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_invalid_arg.py +375 -0
  7. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_numba.py +129 -0
  8. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_apply.py +669 -0
  9. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_apply_relabeling.py +33 -0
  10. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_transform.py +84 -0
  11. .venv/lib/python3.13/site-packages/pandas/tests/apply/test_str.py +307 -0
  12. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/__init__.py +0 -0
  13. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/common.py +158 -0
  14. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/conftest.py +139 -0
  15. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_array_ops.py +78 -0
  16. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_bool.py +28 -0
  17. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_categorical.py +25 -0
  18. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_datetime64.py +2499 -0
  19. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_interval.py +308 -0
  20. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_numeric.py +1588 -0
  21. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_object.py +410 -0
  22. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_period.py +1679 -0
  23. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_string.py +526 -0
  24. .venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_timedelta64.py +2336 -0
  25. .venv/lib/python3.13/site-packages/pandas/tests/arrays/__init__.py +0 -0
  26. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/__init__.py +0 -0
  27. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_arithmetic.py +134 -0
  28. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_astype.py +59 -0
  29. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_comparison.py +60 -0
  30. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_construction.py +335 -0
  31. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_function.py +126 -0
  32. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_indexing.py +13 -0
  33. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_logical.py +255 -0
  34. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_ops.py +27 -0
  35. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_reduction.py +71 -0
  36. .venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_repr.py +13 -0
  37. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/__init__.py +0 -0
  38. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_algos.py +96 -0
  39. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_analytics.py +353 -0
  40. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_api.py +493 -0
  41. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_astype.py +167 -0
  42. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_constructors.py +836 -0
  43. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_dtypes.py +147 -0
  44. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_indexing.py +386 -0
  45. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_map.py +136 -0
  46. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_missing.py +136 -0
  47. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_operators.py +408 -0
  48. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_replace.py +71 -0
  49. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_repr.py +561 -0
  50. .venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_sorting.py +128 -0
.venv/lib/python3.13/site-packages/pandas/tests/apply/common.py ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ from pandas.core.groupby.base import transformation_kernels
2
+
3
+ # There is no Series.cumcount or DataFrame.cumcount
4
+ series_transform_kernels = [
5
+ x for x in sorted(transformation_kernels) if x != "cumcount"
6
+ ]
7
+ frame_transform_kernels = [x for x in sorted(transformation_kernels) if x != "cumcount"]
.venv/lib/python3.13/site-packages/pandas/tests/apply/conftest.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas import (
5
+ DataFrame,
6
+ Series,
7
+ )
8
+ from pandas.api.executors import BaseExecutionEngine
9
+
10
+
11
+ class MockExecutionEngine(BaseExecutionEngine):
12
+ """
13
+ Execution Engine to test if the execution engine interface receives and
14
+ uses all parameters provided by the user.
15
+
16
+ Making this engine work as the default Python engine by calling it, no extra
17
+ functionality is implemented here.
18
+
19
+ When testing, this will be called when this engine is provided, and then the
20
+ same pandas.map and pandas.apply function will be called, but without engine,
21
+ executing the default behavior from the python engine.
22
+ """
23
+
24
+ def map(data, func, args, kwargs, decorator, skip_na):
25
+ kwargs_to_pass = kwargs if isinstance(data, DataFrame) else {}
26
+ return data.map(func, na_action="ignore" if skip_na else None, **kwargs_to_pass)
27
+
28
+ def apply(data, func, args, kwargs, decorator, axis):
29
+ if isinstance(data, Series):
30
+ return data.apply(func, convert_dtype=True, args=args, by_row=False)
31
+ elif isinstance(data, DataFrame):
32
+ return data.apply(
33
+ func,
34
+ axis=axis,
35
+ raw=False,
36
+ result_type=None,
37
+ args=args,
38
+ by_row="compat",
39
+ **kwargs,
40
+ )
41
+ else:
42
+ assert isinstance(data, np.ndarray)
43
+
44
+ def wrap_function(func):
45
+ # https://github.com/numpy/numpy/issues/8352
46
+ def wrapper(*args, **kwargs):
47
+ result = func(*args, **kwargs)
48
+ if isinstance(result, str):
49
+ result = np.array(result, dtype=object)
50
+ return result
51
+
52
+ return wrapper
53
+
54
+ return np.apply_along_axis(wrap_function(func), axis, data, *args, **kwargs)
55
+
56
+
57
+ class MockEngineDecorator:
58
+ __pandas_udf__ = MockExecutionEngine
59
+
60
+
61
+ @pytest.fixture(params=[None, MockEngineDecorator])
62
+ def engine(request):
63
+ return request.param
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_apply.py ADDED
@@ -0,0 +1,1875 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datetime import datetime
2
+ import warnings
3
+
4
+ import numpy as np
5
+ import pytest
6
+
7
+ from pandas.compat import is_platform_arm
8
+
9
+ from pandas.core.dtypes.dtypes import CategoricalDtype
10
+
11
+ import pandas as pd
12
+ from pandas import (
13
+ DataFrame,
14
+ MultiIndex,
15
+ Series,
16
+ Timestamp,
17
+ date_range,
18
+ )
19
+ import pandas._testing as tm
20
+ from pandas.tests.apply.conftest import MockEngineDecorator
21
+ from pandas.tests.frame.common import zip_frames
22
+ from pandas.util.version import Version
23
+
24
+
25
+ @pytest.fixture
26
+ def int_frame_const_col():
27
+ """
28
+ Fixture for DataFrame of ints which are constant per column
29
+
30
+ Columns are ['A', 'B', 'C'], with values (per column): [1, 2, 3]
31
+ """
32
+ df = DataFrame(
33
+ np.tile(np.arange(3, dtype="int64"), 6).reshape(6, -1) + 1,
34
+ columns=["A", "B", "C"],
35
+ )
36
+ return df
37
+
38
+
39
+ @pytest.fixture(
40
+ params=[
41
+ "python",
42
+ pytest.param("numba", marks=pytest.mark.single_cpu),
43
+ MockEngineDecorator,
44
+ ]
45
+ )
46
+ def engine(request):
47
+ if request.param == "numba":
48
+ pytest.importorskip("numba")
49
+ return request.param
50
+
51
+
52
+ def test_apply(float_frame, engine, request):
53
+ if engine == "numba":
54
+ mark = pytest.mark.xfail(reason="numba engine not supporting numpy ufunc yet")
55
+ request.node.add_marker(mark)
56
+ with np.errstate(all="ignore"):
57
+ # ufunc
58
+ result = np.sqrt(float_frame["A"])
59
+ expected = float_frame.apply(np.sqrt, engine=engine)["A"]
60
+ tm.assert_series_equal(result, expected)
61
+
62
+ # aggregator
63
+ result = float_frame.apply(np.mean, engine=engine)["A"]
64
+ expected = np.mean(float_frame["A"])
65
+ assert result == expected
66
+
67
+ d = float_frame.index[0]
68
+ result = float_frame.apply(np.mean, axis=1, engine=engine)
69
+ expected = np.mean(float_frame.xs(d))
70
+ assert result[d] == expected
71
+ assert result.index is float_frame.index
72
+
73
+
74
+ @pytest.mark.parametrize("axis", [0, 1])
75
+ @pytest.mark.parametrize("raw", [True, False])
76
+ @pytest.mark.parametrize("nopython", [True, False])
77
+ def test_apply_args(float_frame, axis, raw, engine, nopython):
78
+ numba = pytest.importorskip("numba")
79
+ if (
80
+ engine == "numba"
81
+ and Version(numba.__version__) == Version("0.61")
82
+ and is_platform_arm()
83
+ ):
84
+ pytest.skip(f"Segfaults on ARM platforms with numba {numba.__version__}")
85
+ engine_kwargs = {"nopython": nopython}
86
+ result = float_frame.apply(
87
+ lambda x, y: x + y,
88
+ axis,
89
+ args=(1,),
90
+ raw=raw,
91
+ engine=engine,
92
+ engine_kwargs=engine_kwargs,
93
+ )
94
+ expected = float_frame + 1
95
+ tm.assert_frame_equal(result, expected)
96
+
97
+ # GH:58712
98
+ result = float_frame.apply(
99
+ lambda x, a, b: x + a + b,
100
+ args=(1,),
101
+ b=2,
102
+ raw=raw,
103
+ engine=engine,
104
+ engine_kwargs=engine_kwargs,
105
+ )
106
+ expected = float_frame + 3
107
+ tm.assert_frame_equal(result, expected)
108
+
109
+ if engine == "numba":
110
+ # py signature binding
111
+ with pytest.raises(TypeError, match="missing a required argument: 'a'"):
112
+ float_frame.apply(
113
+ lambda x, a: x + a,
114
+ b=2,
115
+ raw=raw,
116
+ engine=engine,
117
+ engine_kwargs=engine_kwargs,
118
+ )
119
+
120
+ # keyword-only arguments are not supported in numba
121
+ with pytest.raises(
122
+ pd.errors.NumbaUtilError,
123
+ match="numba does not support keyword-only arguments",
124
+ ):
125
+ float_frame.apply(
126
+ lambda x, a, *, b: x + a + b,
127
+ args=(1,),
128
+ b=2,
129
+ raw=raw,
130
+ engine=engine,
131
+ engine_kwargs=engine_kwargs,
132
+ )
133
+
134
+ with pytest.raises(
135
+ pd.errors.NumbaUtilError,
136
+ match="numba does not support keyword-only arguments",
137
+ ):
138
+ float_frame.apply(
139
+ lambda *x, b: x[0] + x[1] + b,
140
+ args=(1,),
141
+ b=2,
142
+ raw=raw,
143
+ engine=engine,
144
+ engine_kwargs=engine_kwargs,
145
+ )
146
+
147
+
148
+ def test_apply_categorical_func():
149
+ # GH 9573
150
+ df = DataFrame({"c0": ["A", "A", "B", "B"], "c1": ["C", "C", "D", "D"]})
151
+ result = df.apply(lambda ts: ts.astype("category"))
152
+
153
+ assert result.shape == (4, 2)
154
+ assert isinstance(result["c0"].dtype, CategoricalDtype)
155
+ assert isinstance(result["c1"].dtype, CategoricalDtype)
156
+
157
+
158
+ def test_apply_axis1_with_ea():
159
+ # GH#36785
160
+ expected = DataFrame({"A": [Timestamp("2013-01-01", tz="UTC")]})
161
+ result = expected.apply(lambda x: x, axis=1)
162
+ tm.assert_frame_equal(result, expected)
163
+
164
+
165
+ @pytest.mark.parametrize(
166
+ "data, dtype",
167
+ [(1, None), (1, CategoricalDtype([1])), (Timestamp("2013-01-01", tz="UTC"), None)],
168
+ )
169
+ def test_agg_axis1_duplicate_index(data, dtype):
170
+ # GH 42380
171
+ expected = DataFrame([[data], [data]], index=["a", "a"], dtype=dtype)
172
+ result = expected.agg(lambda x: x, axis=1)
173
+ tm.assert_frame_equal(result, expected)
174
+
175
+
176
+ def test_apply_mixed_datetimelike():
177
+ # mixed datetimelike
178
+ # GH 7778
179
+ expected = DataFrame(
180
+ {
181
+ "A": date_range("20130101", periods=3),
182
+ "B": pd.to_timedelta(np.arange(3), unit="s"),
183
+ }
184
+ )
185
+ result = expected.apply(lambda x: x, axis=1)
186
+ tm.assert_frame_equal(result, expected)
187
+
188
+
189
+ @pytest.mark.parametrize("func", [np.sqrt, np.mean])
190
+ def test_apply_empty(func, engine):
191
+ # empty
192
+ empty_frame = DataFrame()
193
+
194
+ result = empty_frame.apply(func, engine=engine)
195
+ assert result.empty
196
+
197
+
198
+ def test_apply_float_frame(float_frame, engine):
199
+ no_rows = float_frame[:0]
200
+ result = no_rows.apply(lambda x: x.mean(), engine=engine)
201
+ expected = Series(np.nan, index=float_frame.columns)
202
+ tm.assert_series_equal(result, expected)
203
+
204
+ no_cols = float_frame.loc[:, []]
205
+ result = no_cols.apply(lambda x: x.mean(), axis=1, engine=engine)
206
+ expected = Series(np.nan, index=float_frame.index)
207
+ tm.assert_series_equal(result, expected)
208
+
209
+
210
+ def test_apply_empty_except_index(engine):
211
+ # GH 2476
212
+ expected = DataFrame(index=["a"])
213
+ result = expected.apply(lambda x: x["a"], axis=1, engine=engine)
214
+ tm.assert_frame_equal(result, expected)
215
+
216
+
217
+ def test_apply_with_reduce_empty():
218
+ # reduce with an empty DataFrame
219
+ empty_frame = DataFrame()
220
+
221
+ x = []
222
+ result = empty_frame.apply(x.append, axis=1, result_type="expand")
223
+ tm.assert_frame_equal(result, empty_frame)
224
+ result = empty_frame.apply(x.append, axis=1, result_type="reduce")
225
+ expected = Series([], dtype=np.float64)
226
+ tm.assert_series_equal(result, expected)
227
+
228
+ empty_with_cols = DataFrame(columns=["a", "b", "c"])
229
+ result = empty_with_cols.apply(x.append, axis=1, result_type="expand")
230
+ tm.assert_frame_equal(result, empty_with_cols)
231
+ result = empty_with_cols.apply(x.append, axis=1, result_type="reduce")
232
+ expected = Series([], dtype=np.float64)
233
+ tm.assert_series_equal(result, expected)
234
+
235
+ # Ensure that x.append hasn't been called
236
+ assert x == []
237
+
238
+
239
+ @pytest.mark.parametrize("func", ["sum", "prod", "any", "all"])
240
+ def test_apply_funcs_over_empty(func):
241
+ # GH 28213
242
+ df = DataFrame(columns=["a", "b", "c"])
243
+
244
+ result = df.apply(getattr(np, func))
245
+ expected = getattr(df, func)()
246
+ if func in ("sum", "prod"):
247
+ expected = expected.astype(float)
248
+ tm.assert_series_equal(result, expected)
249
+
250
+
251
+ def test_nunique_empty():
252
+ # GH 28213
253
+ df = DataFrame(columns=["a", "b", "c"])
254
+
255
+ result = df.nunique()
256
+ expected = Series(0, index=df.columns)
257
+ tm.assert_series_equal(result, expected)
258
+
259
+ result = df.T.nunique()
260
+ expected = Series([], dtype=np.float64)
261
+ tm.assert_series_equal(result, expected)
262
+
263
+
264
+ def test_apply_standard_nonunique():
265
+ df = DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=["a", "a", "c"])
266
+
267
+ result = df.apply(lambda s: s[0], axis=1)
268
+ expected = Series([1, 4, 7], ["a", "a", "c"])
269
+ tm.assert_series_equal(result, expected)
270
+
271
+ result = df.T.apply(lambda s: s[0], axis=0)
272
+ tm.assert_series_equal(result, expected)
273
+
274
+
275
+ def test_apply_broadcast_scalars(float_frame):
276
+ # scalars
277
+ result = float_frame.apply(np.mean, result_type="broadcast")
278
+ expected = DataFrame([float_frame.mean()], index=float_frame.index)
279
+ tm.assert_frame_equal(result, expected)
280
+
281
+
282
+ def test_apply_broadcast_scalars_axis1(float_frame):
283
+ result = float_frame.apply(np.mean, axis=1, result_type="broadcast")
284
+ m = float_frame.mean(axis=1)
285
+ expected = DataFrame(dict.fromkeys(float_frame.columns, m))
286
+ tm.assert_frame_equal(result, expected)
287
+
288
+
289
+ def test_apply_broadcast_lists_columns(float_frame):
290
+ # lists
291
+ result = float_frame.apply(
292
+ lambda x: list(range(len(float_frame.columns))),
293
+ axis=1,
294
+ result_type="broadcast",
295
+ )
296
+ m = list(range(len(float_frame.columns)))
297
+ expected = DataFrame(
298
+ [m] * len(float_frame.index),
299
+ dtype="float64",
300
+ index=float_frame.index,
301
+ columns=float_frame.columns,
302
+ )
303
+ tm.assert_frame_equal(result, expected)
304
+
305
+
306
+ def test_apply_broadcast_lists_index(float_frame):
307
+ result = float_frame.apply(
308
+ lambda x: list(range(len(float_frame.index))), result_type="broadcast"
309
+ )
310
+ m = list(range(len(float_frame.index)))
311
+ expected = DataFrame(
312
+ dict.fromkeys(float_frame.columns, m),
313
+ dtype="float64",
314
+ index=float_frame.index,
315
+ )
316
+ tm.assert_frame_equal(result, expected)
317
+
318
+
319
+ def test_apply_broadcast_list_lambda_func(int_frame_const_col):
320
+ # preserve columns
321
+ df = int_frame_const_col
322
+ result = df.apply(lambda x: [1, 2, 3], axis=1, result_type="broadcast")
323
+ tm.assert_frame_equal(result, df)
324
+
325
+
326
+ def test_apply_broadcast_series_lambda_func(int_frame_const_col):
327
+ df = int_frame_const_col
328
+ result = df.apply(
329
+ lambda x: Series([1, 2, 3], index=list("abc")),
330
+ axis=1,
331
+ result_type="broadcast",
332
+ )
333
+ expected = df.copy()
334
+ tm.assert_frame_equal(result, expected)
335
+
336
+
337
+ @pytest.mark.parametrize("axis", [0, 1])
338
+ def test_apply_raw_float_frame(float_frame, axis, engine):
339
+ if engine == "numba":
340
+ pytest.skip("numba can't handle when UDF returns None.")
341
+
342
+ def _assert_raw(x):
343
+ assert isinstance(x, np.ndarray)
344
+ assert x.ndim == 1
345
+
346
+ float_frame.apply(_assert_raw, axis=axis, engine=engine, raw=True)
347
+
348
+
349
+ @pytest.mark.parametrize("axis", [0, 1])
350
+ def test_apply_raw_float_frame_lambda(float_frame, axis, engine):
351
+ result = float_frame.apply(np.mean, axis=axis, engine=engine, raw=True)
352
+ expected = float_frame.apply(lambda x: x.values.mean(), axis=axis)
353
+ tm.assert_series_equal(result, expected)
354
+
355
+
356
+ def test_apply_raw_float_frame_no_reduction(float_frame, engine):
357
+ # no reduction
358
+ result = float_frame.apply(lambda x: x * 2, engine=engine, raw=True)
359
+ expected = float_frame * 2
360
+ tm.assert_frame_equal(result, expected)
361
+
362
+
363
+ @pytest.mark.parametrize("axis", [0, 1])
364
+ def test_apply_raw_mixed_type_frame(axis, engine):
365
+ if engine == "numba":
366
+ pytest.skip("isinstance check doesn't work with numba")
367
+
368
+ def _assert_raw(x):
369
+ assert isinstance(x, np.ndarray)
370
+ assert x.ndim == 1
371
+
372
+ # Mixed dtype (GH-32423)
373
+ df = DataFrame(
374
+ {
375
+ "a": 1.0,
376
+ "b": 2,
377
+ "c": "foo",
378
+ "float32": np.array([1.0] * 10, dtype="float32"),
379
+ "int32": np.array([1] * 10, dtype="int32"),
380
+ },
381
+ index=np.arange(10),
382
+ )
383
+ df.apply(_assert_raw, axis=axis, engine=engine, raw=True)
384
+
385
+
386
+ def test_apply_axis1(float_frame):
387
+ d = float_frame.index[0]
388
+ result = float_frame.apply(np.mean, axis=1)[d]
389
+ expected = np.mean(float_frame.xs(d))
390
+ assert result == expected
391
+
392
+
393
+ def test_apply_mixed_dtype_corner():
394
+ df = DataFrame({"A": ["foo"], "B": [1.0]})
395
+ result = df[:0].apply(np.mean, axis=1)
396
+ # the result here is actually kind of ambiguous, should it be a Series
397
+ # or a DataFrame?
398
+ expected = Series(dtype=np.float64)
399
+ tm.assert_series_equal(result, expected)
400
+
401
+
402
+ def test_apply_mixed_dtype_corner_indexing():
403
+ df = DataFrame({"A": ["foo"], "B": [1.0]})
404
+ result = df.apply(lambda x: x["A"], axis=1)
405
+ expected = Series(["foo"], index=range(1))
406
+ tm.assert_series_equal(result, expected)
407
+
408
+ result = df.apply(lambda x: x["B"], axis=1)
409
+ expected = Series([1.0], index=range(1))
410
+ tm.assert_series_equal(result, expected)
411
+
412
+
413
+ @pytest.mark.filterwarnings("ignore::RuntimeWarning")
414
+ @pytest.mark.parametrize("ax", ["index", "columns"])
415
+ @pytest.mark.parametrize(
416
+ "func", [lambda x: x, lambda x: x.mean()], ids=["identity", "mean"]
417
+ )
418
+ @pytest.mark.parametrize("raw", [True, False])
419
+ @pytest.mark.parametrize("axis", [0, 1])
420
+ def test_apply_empty_infer_type(ax, func, raw, axis, engine, request):
421
+ df = DataFrame(**{ax: ["a", "b", "c"]})
422
+
423
+ with np.errstate(all="ignore"):
424
+ test_res = func(np.array([], dtype="f8"))
425
+ is_reduction = not isinstance(test_res, np.ndarray)
426
+
427
+ result = df.apply(func, axis=axis, engine=engine, raw=raw)
428
+ if is_reduction:
429
+ agg_axis = df._get_agg_axis(axis)
430
+ assert isinstance(result, Series)
431
+ assert result.index is agg_axis
432
+ else:
433
+ assert isinstance(result, DataFrame)
434
+
435
+
436
+ def test_apply_empty_infer_type_broadcast():
437
+ no_cols = DataFrame(index=["a", "b", "c"])
438
+ result = no_cols.apply(lambda x: x.mean(), result_type="broadcast")
439
+ assert isinstance(result, DataFrame)
440
+
441
+
442
+ def test_apply_with_args_kwds_add_some(float_frame):
443
+ def add_some(x, howmuch=0):
444
+ return x + howmuch
445
+
446
+ result = float_frame.apply(add_some, howmuch=2)
447
+ expected = float_frame.apply(lambda x: x + 2)
448
+ tm.assert_frame_equal(result, expected)
449
+
450
+
451
+ def test_apply_with_args_kwds_agg_and_add(float_frame):
452
+ def agg_and_add(x, howmuch=0):
453
+ return x.mean() + howmuch
454
+
455
+ result = float_frame.apply(agg_and_add, howmuch=2)
456
+ expected = float_frame.apply(lambda x: x.mean() + 2)
457
+ tm.assert_series_equal(result, expected)
458
+
459
+
460
+ def test_apply_with_args_kwds_subtract_and_divide(float_frame):
461
+ def subtract_and_divide(x, sub, divide=1):
462
+ return (x - sub) / divide
463
+
464
+ result = float_frame.apply(subtract_and_divide, args=(2,), divide=2)
465
+ expected = float_frame.apply(lambda x: (x - 2.0) / 2.0)
466
+ tm.assert_frame_equal(result, expected)
467
+
468
+
469
+ def test_apply_yield_list(float_frame):
470
+ result = float_frame.apply(list)
471
+ tm.assert_frame_equal(result, float_frame)
472
+
473
+
474
+ def test_apply_reduce_Series(float_frame):
475
+ float_frame.iloc[::2, float_frame.columns.get_loc("A")] = np.nan
476
+ expected = float_frame.mean(axis=1)
477
+ result = float_frame.apply(np.mean, axis=1)
478
+ tm.assert_series_equal(result, expected)
479
+
480
+
481
+ def test_apply_reduce_to_dict():
482
+ # GH 25196 37544
483
+ data = DataFrame([[1, 2], [3, 4]], columns=["c0", "c1"], index=["i0", "i1"])
484
+
485
+ result = data.apply(dict, axis=0)
486
+ expected = Series([{"i0": 1, "i1": 3}, {"i0": 2, "i1": 4}], index=data.columns)
487
+ tm.assert_series_equal(result, expected)
488
+
489
+ result = data.apply(dict, axis=1)
490
+ expected = Series([{"c0": 1, "c1": 2}, {"c0": 3, "c1": 4}], index=data.index)
491
+ tm.assert_series_equal(result, expected)
492
+
493
+
494
+ def test_apply_differently_indexed():
495
+ df = DataFrame(np.random.default_rng(2).standard_normal((20, 10)))
496
+
497
+ result = df.apply(Series.describe, axis=0)
498
+ expected = DataFrame({i: v.describe() for i, v in df.items()}, columns=df.columns)
499
+ tm.assert_frame_equal(result, expected)
500
+
501
+ result = df.apply(Series.describe, axis=1)
502
+ expected = DataFrame({i: v.describe() for i, v in df.T.items()}, columns=df.index).T
503
+ tm.assert_frame_equal(result, expected)
504
+
505
+
506
+ def test_apply_bug():
507
+ # GH 6125
508
+ positions = DataFrame(
509
+ [
510
+ [1, "ABC0", 50],
511
+ [1, "YUM0", 20],
512
+ [1, "DEF0", 20],
513
+ [2, "ABC1", 50],
514
+ [2, "YUM1", 20],
515
+ [2, "DEF1", 20],
516
+ ],
517
+ columns=["a", "market", "position"],
518
+ )
519
+
520
+ def f(r):
521
+ return r["market"]
522
+
523
+ expected = positions.apply(f, axis=1)
524
+
525
+ positions = DataFrame(
526
+ [
527
+ [datetime(2013, 1, 1), "ABC0", 50],
528
+ [datetime(2013, 1, 2), "YUM0", 20],
529
+ [datetime(2013, 1, 3), "DEF0", 20],
530
+ [datetime(2013, 1, 4), "ABC1", 50],
531
+ [datetime(2013, 1, 5), "YUM1", 20],
532
+ [datetime(2013, 1, 6), "DEF1", 20],
533
+ ],
534
+ columns=["a", "market", "position"],
535
+ )
536
+ result = positions.apply(f, axis=1)
537
+ tm.assert_series_equal(result, expected)
538
+
539
+
540
+ def test_apply_convert_objects():
541
+ expected = DataFrame(
542
+ {
543
+ "A": [
544
+ "foo",
545
+ "foo",
546
+ "foo",
547
+ "foo",
548
+ "bar",
549
+ "bar",
550
+ "bar",
551
+ "bar",
552
+ "foo",
553
+ "foo",
554
+ "foo",
555
+ ],
556
+ "B": [
557
+ "one",
558
+ "one",
559
+ "one",
560
+ "two",
561
+ "one",
562
+ "one",
563
+ "one",
564
+ "two",
565
+ "two",
566
+ "two",
567
+ "one",
568
+ ],
569
+ "C": [
570
+ "dull",
571
+ "dull",
572
+ "shiny",
573
+ "dull",
574
+ "dull",
575
+ "shiny",
576
+ "shiny",
577
+ "dull",
578
+ "shiny",
579
+ "shiny",
580
+ "shiny",
581
+ ],
582
+ "D": np.random.default_rng(2).standard_normal(11),
583
+ "E": np.random.default_rng(2).standard_normal(11),
584
+ "F": np.random.default_rng(2).standard_normal(11),
585
+ }
586
+ )
587
+
588
+ result = expected.apply(lambda x: x, axis=1)
589
+ tm.assert_frame_equal(result, expected)
590
+
591
+
592
+ def test_apply_attach_name(float_frame):
593
+ result = float_frame.apply(lambda x: x.name)
594
+ expected = Series(float_frame.columns, index=float_frame.columns)
595
+ tm.assert_series_equal(result, expected)
596
+
597
+
598
+ def test_apply_attach_name_axis1(float_frame):
599
+ result = float_frame.apply(lambda x: x.name, axis=1)
600
+ expected = Series(float_frame.index, index=float_frame.index)
601
+ tm.assert_series_equal(result, expected)
602
+
603
+
604
+ def test_apply_attach_name_non_reduction(float_frame):
605
+ # non-reductions
606
+ result = float_frame.apply(lambda x: np.repeat(x.name, len(x)))
607
+ expected = DataFrame(
608
+ np.tile(float_frame.columns, (len(float_frame.index), 1)),
609
+ index=float_frame.index,
610
+ columns=float_frame.columns,
611
+ )
612
+ tm.assert_frame_equal(result, expected)
613
+
614
+
615
+ def test_apply_attach_name_non_reduction_axis1(float_frame):
616
+ result = float_frame.apply(lambda x: np.repeat(x.name, len(x)), axis=1)
617
+ expected = Series(
618
+ np.repeat(t[0], len(float_frame.columns)) for t in float_frame.itertuples()
619
+ )
620
+ expected.index = float_frame.index
621
+ tm.assert_series_equal(result, expected)
622
+
623
+
624
+ def test_apply_multi_index():
625
+ index = MultiIndex.from_arrays([["a", "a", "b"], ["c", "d", "d"]])
626
+ s = DataFrame([[1, 2], [3, 4], [5, 6]], index=index, columns=["col1", "col2"])
627
+ result = s.apply(lambda x: Series({"min": min(x), "max": max(x)}), 1)
628
+ expected = DataFrame([[1, 2], [3, 4], [5, 6]], index=index, columns=["min", "max"])
629
+ tm.assert_frame_equal(result, expected, check_like=True)
630
+
631
+
632
+ @pytest.mark.parametrize(
633
+ "df, dicts",
634
+ [
635
+ [
636
+ DataFrame([["foo", "bar"], ["spam", "eggs"]]),
637
+ Series([{0: "foo", 1: "spam"}, {0: "bar", 1: "eggs"}]),
638
+ ],
639
+ [DataFrame([[0, 1], [2, 3]]), Series([{0: 0, 1: 2}, {0: 1, 1: 3}])],
640
+ ],
641
+ )
642
+ def test_apply_dict(df, dicts):
643
+ # GH 8735
644
+ fn = lambda x: x.to_dict()
645
+ reduce_true = df.apply(fn, result_type="reduce")
646
+ reduce_false = df.apply(fn, result_type="expand")
647
+ reduce_none = df.apply(fn)
648
+
649
+ tm.assert_series_equal(reduce_true, dicts)
650
+ tm.assert_frame_equal(reduce_false, df)
651
+ tm.assert_series_equal(reduce_none, dicts)
652
+
653
+
654
+ def test_apply_non_numpy_dtype():
655
+ # GH 12244
656
+ df = DataFrame({"dt": date_range("2015-01-01", periods=3, tz="Europe/Brussels")})
657
+ result = df.apply(lambda x: x)
658
+ tm.assert_frame_equal(result, df)
659
+
660
+ result = df.apply(lambda x: x + pd.Timedelta("1day"))
661
+ expected = DataFrame(
662
+ {"dt": date_range("2015-01-02", periods=3, tz="Europe/Brussels")}
663
+ )
664
+ tm.assert_frame_equal(result, expected)
665
+
666
+
667
+ def test_apply_non_numpy_dtype_category():
668
+ df = DataFrame({"dt": ["a", "b", "c", "a"]}, dtype="category")
669
+ result = df.apply(lambda x: x)
670
+ tm.assert_frame_equal(result, df)
671
+
672
+
673
+ def test_apply_dup_names_multi_agg():
674
+ # GH 21063
675
+ df = DataFrame([[0, 1], [2, 3]], columns=["a", "a"])
676
+ expected = DataFrame([[0, 1]], columns=["a", "a"], index=["min"])
677
+ result = df.agg(["min"])
678
+
679
+ tm.assert_frame_equal(result, expected)
680
+
681
+
682
+ @pytest.mark.parametrize("op", ["apply", "agg"])
683
+ def test_apply_nested_result_axis_1(op):
684
+ # GH 13820
685
+ def apply_list(row):
686
+ return [2 * row["A"], 2 * row["C"], 2 * row["B"]]
687
+
688
+ df = DataFrame(np.zeros((4, 4)), columns=list("ABCD"))
689
+ result = getattr(df, op)(apply_list, axis=1)
690
+ expected = Series(
691
+ [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
692
+ )
693
+ tm.assert_series_equal(result, expected)
694
+
695
+
696
+ def test_apply_noreduction_tzaware_object():
697
+ # https://github.com/pandas-dev/pandas/issues/31505
698
+ expected = DataFrame(
699
+ {"foo": [Timestamp("2020", tz="UTC")]}, dtype="datetime64[ns, UTC]"
700
+ )
701
+ result = expected.apply(lambda x: x)
702
+ tm.assert_frame_equal(result, expected)
703
+ result = expected.apply(lambda x: x.copy())
704
+ tm.assert_frame_equal(result, expected)
705
+
706
+
707
+ def test_apply_function_runs_once():
708
+ # https://github.com/pandas-dev/pandas/issues/30815
709
+
710
+ df = DataFrame({"a": [1, 2, 3]})
711
+ names = [] # Save row names function is applied to
712
+
713
+ def reducing_function(row):
714
+ names.append(row.name)
715
+
716
+ def non_reducing_function(row):
717
+ names.append(row.name)
718
+ return row
719
+
720
+ for func in [reducing_function, non_reducing_function]:
721
+ del names[:]
722
+
723
+ df.apply(func, axis=1)
724
+ assert names == list(df.index)
725
+
726
+
727
+ def test_apply_raw_function_runs_once(engine):
728
+ # https://github.com/pandas-dev/pandas/issues/34506
729
+ if engine == "numba":
730
+ pytest.skip("appending to list outside of numba func is not supported")
731
+
732
+ df = DataFrame({"a": [1, 2, 3]})
733
+ values = [] # Save row values function is applied to
734
+
735
+ def reducing_function(row):
736
+ values.extend(row)
737
+
738
+ def non_reducing_function(row):
739
+ values.extend(row)
740
+ return row
741
+
742
+ for func in [reducing_function, non_reducing_function]:
743
+ del values[:]
744
+
745
+ df.apply(func, engine=engine, raw=True, axis=1)
746
+ assert values == list(df.a.to_list())
747
+
748
+
749
+ def test_apply_with_byte_string():
750
+ # GH 34529
751
+ df = DataFrame(np.array([b"abcd", b"efgh"]), columns=["col"])
752
+ expected = DataFrame(np.array([b"abcd", b"efgh"]), columns=["col"], dtype=object)
753
+ # After we make the apply we expect a dataframe just
754
+ # like the original but with the object datatype
755
+ result = df.apply(lambda x: x.astype("object"))
756
+ tm.assert_frame_equal(result, expected)
757
+
758
+
759
+ @pytest.mark.parametrize("val", ["asd", 12, None, np.nan])
760
+ def test_apply_category_equalness(val):
761
+ # Check if categorical comparisons on apply, GH 21239
762
+ df_values = ["asd", None, 12, "asd", "cde", np.nan]
763
+ df = DataFrame({"a": df_values}, dtype="category")
764
+
765
+ result = df.a.apply(lambda x: x == val)
766
+ expected = Series(
767
+ [False if pd.isnull(x) else x == val for x in df_values], name="a"
768
+ )
769
+ # False since behavior of NaN for categorical dtype has been changed (GH 59966)
770
+ tm.assert_series_equal(result, expected)
771
+
772
+
773
+ # the user has supplied an opaque UDF where
774
+ # they are transforming the input that requires
775
+ # us to infer the output
776
+
777
+
778
+ def test_infer_row_shape():
779
+ # GH 17437
780
+ # if row shape is changing, infer it
781
+ df = DataFrame(np.random.default_rng(2).random((10, 2)))
782
+ result = df.apply(np.fft.fft, axis=0).shape
783
+ assert result == (10, 2)
784
+
785
+ result = df.apply(np.fft.rfft, axis=0).shape
786
+ assert result == (6, 2)
787
+
788
+
789
+ @pytest.mark.parametrize(
790
+ "ops, by_row, expected",
791
+ [
792
+ ({"a": lambda x: x + 1}, "compat", DataFrame({"a": [2, 3]})),
793
+ ({"a": lambda x: x + 1}, False, DataFrame({"a": [2, 3]})),
794
+ ({"a": lambda x: x.sum()}, "compat", Series({"a": 3})),
795
+ ({"a": lambda x: x.sum()}, False, Series({"a": 3})),
796
+ (
797
+ {"a": ["sum", np.sum, lambda x: x.sum()]},
798
+ "compat",
799
+ DataFrame({"a": [3, 3, 3]}, index=["sum", "sum", "<lambda>"]),
800
+ ),
801
+ (
802
+ {"a": ["sum", np.sum, lambda x: x.sum()]},
803
+ False,
804
+ DataFrame({"a": [3, 3, 3]}, index=["sum", "sum", "<lambda>"]),
805
+ ),
806
+ ({"a": lambda x: 1}, "compat", DataFrame({"a": [1, 1]})),
807
+ ({"a": lambda x: 1}, False, Series({"a": 1})),
808
+ ],
809
+ )
810
+ def test_dictlike_lambda(ops, by_row, expected):
811
+ # GH53601
812
+ df = DataFrame({"a": [1, 2]})
813
+ result = df.apply(ops, by_row=by_row)
814
+ tm.assert_equal(result, expected)
815
+
816
+
817
+ @pytest.mark.parametrize(
818
+ "ops",
819
+ [
820
+ {"a": lambda x: x + 1},
821
+ {"a": lambda x: x.sum()},
822
+ {"a": ["sum", np.sum, lambda x: x.sum()]},
823
+ {"a": lambda x: 1},
824
+ ],
825
+ )
826
+ def test_dictlike_lambda_raises(ops):
827
+ # GH53601
828
+ df = DataFrame({"a": [1, 2]})
829
+ with pytest.raises(ValueError, match="by_row=True not allowed"):
830
+ df.apply(ops, by_row=True)
831
+
832
+
833
+ def test_with_dictlike_columns():
834
+ # GH 17602
835
+ df = DataFrame([[1, 2], [1, 2]], columns=["a", "b"])
836
+ result = df.apply(lambda x: {"s": x["a"] + x["b"]}, axis=1)
837
+ expected = Series([{"s": 3} for t in df.itertuples()])
838
+ tm.assert_series_equal(result, expected)
839
+
840
+ df["tm"] = [
841
+ Timestamp("2017-05-01 00:00:00"),
842
+ Timestamp("2017-05-02 00:00:00"),
843
+ ]
844
+ result = df.apply(lambda x: {"s": x["a"] + x["b"]}, axis=1)
845
+ tm.assert_series_equal(result, expected)
846
+
847
+ # compose a series
848
+ result = (df["a"] + df["b"]).apply(lambda x: {"s": x})
849
+ expected = Series([{"s": 3}, {"s": 3}])
850
+ tm.assert_series_equal(result, expected)
851
+
852
+
853
+ def test_with_dictlike_columns_with_datetime():
854
+ # GH 18775
855
+ df = DataFrame()
856
+ df["author"] = ["X", "Y", "Z"]
857
+ df["publisher"] = ["BBC", "NBC", "N24"]
858
+ df["date"] = pd.to_datetime(
859
+ ["17-10-2010 07:15:30", "13-05-2011 08:20:35", "15-01-2013 09:09:09"],
860
+ dayfirst=True,
861
+ )
862
+ result = df.apply(lambda x: {}, axis=1)
863
+ expected = Series([{}, {}, {}])
864
+ tm.assert_series_equal(result, expected)
865
+
866
+
867
+ def test_with_dictlike_columns_with_infer():
868
+ # GH 17602
869
+ df = DataFrame([[1, 2], [1, 2]], columns=["a", "b"])
870
+ result = df.apply(lambda x: {"s": x["a"] + x["b"]}, axis=1, result_type="expand")
871
+ expected = DataFrame({"s": [3, 3]})
872
+ tm.assert_frame_equal(result, expected)
873
+
874
+ df["tm"] = [
875
+ Timestamp("2017-05-01 00:00:00"),
876
+ Timestamp("2017-05-02 00:00:00"),
877
+ ]
878
+ result = df.apply(lambda x: {"s": x["a"] + x["b"]}, axis=1, result_type="expand")
879
+ tm.assert_frame_equal(result, expected)
880
+
881
+
882
+ @pytest.mark.parametrize(
883
+ "ops, by_row, expected",
884
+ [
885
+ ([lambda x: x + 1], "compat", DataFrame({("a", "<lambda>"): [2, 3]})),
886
+ ([lambda x: x + 1], False, DataFrame({("a", "<lambda>"): [2, 3]})),
887
+ ([lambda x: x.sum()], "compat", DataFrame({"a": [3]}, index=["<lambda>"])),
888
+ ([lambda x: x.sum()], False, DataFrame({"a": [3]}, index=["<lambda>"])),
889
+ (
890
+ ["sum", np.sum, lambda x: x.sum()],
891
+ "compat",
892
+ DataFrame({"a": [3, 3, 3]}, index=["sum", "sum", "<lambda>"]),
893
+ ),
894
+ (
895
+ ["sum", np.sum, lambda x: x.sum()],
896
+ False,
897
+ DataFrame({"a": [3, 3, 3]}, index=["sum", "sum", "<lambda>"]),
898
+ ),
899
+ (
900
+ [lambda x: x + 1, lambda x: 3],
901
+ "compat",
902
+ DataFrame([[2, 3], [3, 3]], columns=[["a", "a"], ["<lambda>", "<lambda>"]]),
903
+ ),
904
+ (
905
+ [lambda x: 2, lambda x: 3],
906
+ False,
907
+ DataFrame({"a": [2, 3]}, ["<lambda>", "<lambda>"]),
908
+ ),
909
+ ],
910
+ )
911
+ def test_listlike_lambda(ops, by_row, expected):
912
+ # GH53601
913
+ df = DataFrame({"a": [1, 2]})
914
+ result = df.apply(ops, by_row=by_row)
915
+ tm.assert_equal(result, expected)
916
+
917
+
918
+ def test_listlike_datetime_index_unsorted():
919
+ # https://github.com/pandas-dev/pandas/pull/62843
920
+ values = [datetime(2024, 1, 1), datetime(2024, 1, 2), datetime(2024, 1, 3)]
921
+ df = DataFrame({"a": [1, 2]}, index=[values[1], values[0]])
922
+ result = df.apply([lambda x: x, lambda x: x.shift(freq="D")], by_row=False)
923
+ expected = DataFrame(
924
+ [[1.0, 2.0], [2.0, np.nan], [np.nan, 1.0]],
925
+ index=[values[1], values[0], values[2]],
926
+ columns=MultiIndex([["a"], ["<lambda>"]], codes=[[0, 0], [0, 0]]),
927
+ )
928
+ tm.assert_frame_equal(result, expected)
929
+
930
+
931
+ def test_dictlike_datetime_index_unsorted():
932
+ # https://github.com/pandas-dev/pandas/pull/62843
933
+ values = [datetime(2024, 1, 1), datetime(2024, 1, 2), datetime(2024, 1, 3)]
934
+ df = DataFrame({"a": [1, 2], "b": [3, 4]}, index=[values[1], values[0]])
935
+ result = df.apply(
936
+ {"a": lambda x: x, "b": lambda x: x.shift(freq="D")}, by_row=False
937
+ )
938
+ expected = DataFrame(
939
+ {
940
+ "a": [1.0, 2.0, np.nan],
941
+ "b": [4.0, np.nan, 3.0],
942
+ },
943
+ index=[values[1], values[0], values[2]],
944
+ )
945
+ tm.assert_frame_equal(result, expected)
946
+
947
+
948
+ @pytest.mark.parametrize(
949
+ "ops",
950
+ [
951
+ [lambda x: x + 1],
952
+ [lambda x: x.sum()],
953
+ ["sum", np.sum, lambda x: x.sum()],
954
+ [lambda x: x + 1, lambda x: 3],
955
+ ],
956
+ )
957
+ def test_listlike_lambda_raises(ops):
958
+ # GH53601
959
+ df = DataFrame({"a": [1, 2]})
960
+ with pytest.raises(ValueError, match="by_row=True not allowed"):
961
+ df.apply(ops, by_row=True)
962
+
963
+
964
+ def test_with_listlike_columns():
965
+ # GH 17348
966
+ df = DataFrame(
967
+ {
968
+ "a": Series(np.random.default_rng(2).standard_normal(4)),
969
+ "b": ["a", "list", "of", "words"],
970
+ "ts": date_range("2016-10-01", periods=4, freq="h"),
971
+ }
972
+ )
973
+
974
+ result = df[["a", "b"]].apply(tuple, axis=1)
975
+ expected = Series([t[1:] for t in df[["a", "b"]].itertuples()])
976
+ tm.assert_series_equal(result, expected)
977
+
978
+ result = df[["a", "ts"]].apply(tuple, axis=1)
979
+ expected = Series([t[1:] for t in df[["a", "ts"]].itertuples()])
980
+ tm.assert_series_equal(result, expected)
981
+
982
+
983
+ def test_with_listlike_columns_returning_list():
984
+ # GH 18919
985
+ df = DataFrame({"x": Series([["a", "b"], ["q"]]), "y": Series([["z"], ["q", "t"]])})
986
+ df.index = MultiIndex.from_tuples([("i0", "j0"), ("i1", "j1")])
987
+
988
+ result = df.apply(lambda row: [el for el in row["x"] if el in row["y"]], axis=1)
989
+ expected = Series([[], ["q"]], index=df.index)
990
+ tm.assert_series_equal(result, expected)
991
+
992
+
993
+ def test_infer_output_shape_columns():
994
+ # GH 18573
995
+
996
+ df = DataFrame(
997
+ {
998
+ "number": [1.0, 2.0],
999
+ "string": ["foo", "bar"],
1000
+ "datetime": [
1001
+ Timestamp("2017-11-29 03:30:00"),
1002
+ Timestamp("2017-11-29 03:45:00"),
1003
+ ],
1004
+ }
1005
+ )
1006
+ result = df.apply(lambda row: (row.number, row.string), axis=1)
1007
+ expected = Series([(t.number, t.string) for t in df.itertuples()])
1008
+ tm.assert_series_equal(result, expected)
1009
+
1010
+
1011
+ def test_infer_output_shape_listlike_columns():
1012
+ # GH 16353
1013
+
1014
+ df = DataFrame(
1015
+ np.random.default_rng(2).standard_normal((6, 3)), columns=["A", "B", "C"]
1016
+ )
1017
+
1018
+ result = df.apply(lambda x: [1, 2, 3], axis=1)
1019
+ expected = Series([[1, 2, 3] for t in df.itertuples()])
1020
+ tm.assert_series_equal(result, expected)
1021
+
1022
+ result = df.apply(lambda x: [1, 2], axis=1)
1023
+ expected = Series([[1, 2] for t in df.itertuples()])
1024
+ tm.assert_series_equal(result, expected)
1025
+
1026
+
1027
+ @pytest.mark.parametrize("val", [1, 2])
1028
+ def test_infer_output_shape_listlike_columns_np_func(val):
1029
+ # GH 17970
1030
+ df = DataFrame({"a": [1, 2, 3]}, index=list("abc"))
1031
+
1032
+ result = df.apply(lambda row: np.ones(val), axis=1)
1033
+ expected = Series([np.ones(val) for t in df.itertuples()], index=df.index)
1034
+ tm.assert_series_equal(result, expected)
1035
+
1036
+
1037
+ def test_infer_output_shape_listlike_columns_with_timestamp():
1038
+ # GH 17892
1039
+ df = DataFrame(
1040
+ {
1041
+ "a": [
1042
+ Timestamp("2010-02-01"),
1043
+ Timestamp("2010-02-04"),
1044
+ Timestamp("2010-02-05"),
1045
+ Timestamp("2010-02-06"),
1046
+ ],
1047
+ "b": [9, 5, 4, 3],
1048
+ "c": [5, 3, 4, 2],
1049
+ "d": [1, 2, 3, 4],
1050
+ }
1051
+ )
1052
+
1053
+ def fun(x):
1054
+ return (1, 2)
1055
+
1056
+ result = df.apply(fun, axis=1)
1057
+ expected = Series([(1, 2) for t in df.itertuples()])
1058
+ tm.assert_series_equal(result, expected)
1059
+
1060
+
1061
+ @pytest.mark.parametrize("lst", [[1, 2, 3], [1, 2]])
1062
+ def test_consistent_coerce_for_shapes(lst):
1063
+ # we want column names to NOT be propagated
1064
+ # just because the shape matches the input shape
1065
+ df = DataFrame(
1066
+ np.random.default_rng(2).standard_normal((4, 3)), columns=["A", "B", "C"]
1067
+ )
1068
+
1069
+ result = df.apply(lambda x: lst, axis=1)
1070
+ expected = Series([lst for t in df.itertuples()])
1071
+ tm.assert_series_equal(result, expected)
1072
+
1073
+
1074
+ def test_consistent_names(int_frame_const_col):
1075
+ # if a Series is returned, we should use the resulting index names
1076
+ df = int_frame_const_col
1077
+
1078
+ result = df.apply(
1079
+ lambda x: Series([1, 2, 3], index=["test", "other", "cols"]), axis=1
1080
+ )
1081
+ expected = int_frame_const_col.rename(
1082
+ columns={"A": "test", "B": "other", "C": "cols"}
1083
+ )
1084
+ tm.assert_frame_equal(result, expected)
1085
+
1086
+ result = df.apply(lambda x: Series([1, 2], index=["test", "other"]), axis=1)
1087
+ expected = expected[["test", "other"]]
1088
+ tm.assert_frame_equal(result, expected)
1089
+
1090
+
1091
+ def test_result_type(int_frame_const_col):
1092
+ # result_type should be consistent no matter which
1093
+ # path we take in the code
1094
+ df = int_frame_const_col
1095
+
1096
+ result = df.apply(lambda x: [1, 2, 3], axis=1, result_type="expand")
1097
+ expected = df.copy()
1098
+ expected.columns = range(3)
1099
+ tm.assert_frame_equal(result, expected)
1100
+
1101
+
1102
+ def test_result_type_shorter_list(int_frame_const_col):
1103
+ # result_type should be consistent no matter which
1104
+ # path we take in the code
1105
+ df = int_frame_const_col
1106
+ result = df.apply(lambda x: [1, 2], axis=1, result_type="expand")
1107
+ expected = df[["A", "B"]].copy()
1108
+ expected.columns = range(2)
1109
+ tm.assert_frame_equal(result, expected)
1110
+
1111
+
1112
+ def test_result_type_broadcast(int_frame_const_col, request, engine):
1113
+ # result_type should be consistent no matter which
1114
+ # path we take in the code
1115
+ if engine == "numba":
1116
+ mark = pytest.mark.xfail(reason="numba engine doesn't support list return")
1117
+ request.node.add_marker(mark)
1118
+ df = int_frame_const_col
1119
+ if engine is MockEngineDecorator:
1120
+ with pytest.raises(
1121
+ NotImplementedError,
1122
+ match="result_type='broadcast' only implemented for the default engine",
1123
+ ):
1124
+ df.apply(
1125
+ lambda x: [1, 2, 3], axis=1, result_type="broadcast", engine=engine
1126
+ )
1127
+ else:
1128
+ # broadcast result
1129
+ result = df.apply(
1130
+ lambda x: [1, 2, 3], axis=1, result_type="broadcast", engine=engine
1131
+ )
1132
+ expected = df.copy()
1133
+ tm.assert_frame_equal(result, expected)
1134
+
1135
+
1136
+ def test_result_type_broadcast_series_func(int_frame_const_col, engine, request):
1137
+ # result_type should be consistent no matter which
1138
+ # path we take in the code
1139
+ if engine == "numba":
1140
+ mark = pytest.mark.xfail(
1141
+ reason="numba Series constructor only support ndarrays not list data"
1142
+ )
1143
+ request.node.add_marker(mark)
1144
+ df = int_frame_const_col
1145
+ columns = ["other", "col", "names"]
1146
+
1147
+ if engine is MockEngineDecorator:
1148
+ with pytest.raises(
1149
+ NotImplementedError,
1150
+ match="result_type='broadcast' only implemented for the default engine",
1151
+ ):
1152
+ df.apply(
1153
+ lambda x: Series([1, 2, 3], index=columns),
1154
+ axis=1,
1155
+ result_type="broadcast",
1156
+ engine=engine,
1157
+ )
1158
+ else:
1159
+ result = df.apply(
1160
+ lambda x: Series([1, 2, 3], index=columns),
1161
+ axis=1,
1162
+ result_type="broadcast",
1163
+ engine=engine,
1164
+ )
1165
+ expected = df.copy()
1166
+ tm.assert_frame_equal(result, expected)
1167
+
1168
+
1169
+ def test_result_type_series_result(int_frame_const_col, engine, request):
1170
+ # result_type should be consistent no matter which
1171
+ # path we take in the code
1172
+ if engine == "numba":
1173
+ mark = pytest.mark.xfail(
1174
+ reason="numba Series constructor only support ndarrays not list data"
1175
+ )
1176
+ request.node.add_marker(mark)
1177
+ df = int_frame_const_col
1178
+ # series result
1179
+ result = df.apply(lambda x: Series([1, 2, 3], index=x.index), axis=1, engine=engine)
1180
+ expected = df.copy()
1181
+ tm.assert_frame_equal(result, expected)
1182
+
1183
+
1184
+ def test_result_type_series_result_other_index(int_frame_const_col, engine, request):
1185
+ # result_type should be consistent no matter which
1186
+ # path we take in the code
1187
+
1188
+ if engine == "numba":
1189
+ mark = pytest.mark.xfail(
1190
+ reason="no support in numba Series constructor for list of columns"
1191
+ )
1192
+ request.node.add_marker(mark)
1193
+ df = int_frame_const_col
1194
+ # series result with other index
1195
+ columns = ["other", "col", "names"]
1196
+ result = df.apply(lambda x: Series([1, 2, 3], index=columns), axis=1, engine=engine)
1197
+ expected = df.copy()
1198
+ expected.columns = columns
1199
+ tm.assert_frame_equal(result, expected)
1200
+
1201
+
1202
+ @pytest.mark.parametrize(
1203
+ "box",
1204
+ [lambda x: list(x), lambda x: tuple(x), lambda x: np.array(x, dtype="int64")],
1205
+ ids=["list", "tuple", "array"],
1206
+ )
1207
+ def test_consistency_for_boxed(box, int_frame_const_col):
1208
+ # passing an array or list should not affect the output shape
1209
+ df = int_frame_const_col
1210
+
1211
+ result = df.apply(lambda x: box([1, 2]), axis=1)
1212
+ expected = Series([box([1, 2]) for t in df.itertuples()])
1213
+ tm.assert_series_equal(result, expected)
1214
+
1215
+ result = df.apply(lambda x: box([1, 2]), axis=1, result_type="expand")
1216
+ expected = int_frame_const_col[["A", "B"]].rename(columns={"A": 0, "B": 1})
1217
+ tm.assert_frame_equal(result, expected)
1218
+
1219
+
1220
+ def test_agg_transform(axis, float_frame):
1221
+ other_axis = 1 if axis in {0, "index"} else 0
1222
+
1223
+ with np.errstate(all="ignore"):
1224
+ f_abs = np.abs(float_frame)
1225
+ f_sqrt = np.sqrt(float_frame)
1226
+
1227
+ # ufunc
1228
+ expected = f_sqrt.copy()
1229
+ result = float_frame.apply(np.sqrt, axis=axis)
1230
+ tm.assert_frame_equal(result, expected)
1231
+
1232
+ # list-like
1233
+ result = float_frame.apply([np.sqrt], axis=axis)
1234
+ expected = f_sqrt.copy()
1235
+ if axis in {0, "index"}:
1236
+ expected.columns = MultiIndex.from_product([float_frame.columns, ["sqrt"]])
1237
+ else:
1238
+ expected.index = MultiIndex.from_product([float_frame.index, ["sqrt"]])
1239
+ tm.assert_frame_equal(result, expected)
1240
+
1241
+ # multiple items in list
1242
+ # these are in the order as if we are applying both
1243
+ # functions per series and then concatting
1244
+ result = float_frame.apply([np.abs, np.sqrt], axis=axis)
1245
+ expected = zip_frames([f_abs, f_sqrt], axis=other_axis)
1246
+ if axis in {0, "index"}:
1247
+ expected.columns = MultiIndex.from_product(
1248
+ [float_frame.columns, ["absolute", "sqrt"]]
1249
+ )
1250
+ else:
1251
+ expected.index = MultiIndex.from_product(
1252
+ [float_frame.index, ["absolute", "sqrt"]]
1253
+ )
1254
+ tm.assert_frame_equal(result, expected)
1255
+
1256
+
1257
+ def test_demo():
1258
+ # demonstration tests
1259
+ df = DataFrame({"A": range(5), "B": 5})
1260
+
1261
+ result = df.agg(["min", "max"])
1262
+ expected = DataFrame(
1263
+ {"A": [0, 4], "B": [5, 5]}, columns=["A", "B"], index=["min", "max"]
1264
+ )
1265
+ tm.assert_frame_equal(result, expected)
1266
+
1267
+
1268
+ def test_demo_dict_agg():
1269
+ # demonstration tests
1270
+ df = DataFrame({"A": range(5), "B": 5})
1271
+ result = df.agg({"A": ["min", "max"], "B": ["sum", "max"]})
1272
+ expected = DataFrame(
1273
+ {"A": [4.0, 0.0, np.nan], "B": [5.0, np.nan, 25.0]},
1274
+ columns=["A", "B"],
1275
+ index=["max", "min", "sum"],
1276
+ )
1277
+ tm.assert_frame_equal(result.reindex_like(expected), expected)
1278
+
1279
+
1280
+ def test_agg_with_name_as_column_name():
1281
+ # GH 36212 - Column name is "name"
1282
+ data = {"name": ["foo", "bar"]}
1283
+ df = DataFrame(data)
1284
+
1285
+ # result's name should be None
1286
+ result = df.agg({"name": "count"})
1287
+ expected = Series({"name": 2})
1288
+ tm.assert_series_equal(result, expected)
1289
+
1290
+ # Check if name is still preserved when aggregating series instead
1291
+ result = df["name"].agg({"name": "count"})
1292
+ expected = Series({"name": 2}, name="name")
1293
+ tm.assert_series_equal(result, expected)
1294
+
1295
+
1296
+ def test_agg_multiple_mixed():
1297
+ # GH 20909
1298
+ mdf = DataFrame(
1299
+ {
1300
+ "A": [1, 2, 3],
1301
+ "B": [1.0, 2.0, 3.0],
1302
+ "C": ["foo", "bar", "baz"],
1303
+ }
1304
+ )
1305
+ expected = DataFrame(
1306
+ {
1307
+ "A": [1, 6],
1308
+ "B": [1.0, 6.0],
1309
+ "C": ["bar", "foobarbaz"],
1310
+ },
1311
+ index=["min", "sum"],
1312
+ )
1313
+ # sorted index
1314
+ result = mdf.agg(["min", "sum"])
1315
+ tm.assert_frame_equal(result, expected)
1316
+
1317
+ result = mdf[["C", "B", "A"]].agg(["sum", "min"])
1318
+ # GH40420: the result of .agg should have an index that is sorted
1319
+ # according to the arguments provided to agg.
1320
+ expected = expected[["C", "B", "A"]].reindex(["sum", "min"])
1321
+ tm.assert_frame_equal(result, expected)
1322
+
1323
+
1324
+ def test_agg_multiple_mixed_raises():
1325
+ # GH 20909
1326
+ mdf = DataFrame(
1327
+ {
1328
+ "A": [1, 2, 3],
1329
+ "B": [1.0, 2.0, 3.0],
1330
+ "C": ["foo", "bar", "baz"],
1331
+ "D": date_range("20130101", periods=3),
1332
+ }
1333
+ )
1334
+
1335
+ # sorted index
1336
+ msg = "does not support operation"
1337
+ with pytest.raises(TypeError, match=msg):
1338
+ mdf.agg(["min", "sum"])
1339
+
1340
+ with pytest.raises(TypeError, match=msg):
1341
+ mdf[["D", "C", "B", "A"]].agg(["sum", "min"])
1342
+
1343
+
1344
+ def test_agg_reduce(axis, float_frame):
1345
+ other_axis = 1 if axis in {0, "index"} else 0
1346
+ name1, name2 = float_frame.axes[other_axis].unique()[:2].sort_values()
1347
+
1348
+ # all reducers
1349
+ expected = pd.concat(
1350
+ [
1351
+ float_frame.mean(axis=axis),
1352
+ float_frame.max(axis=axis),
1353
+ float_frame.sum(axis=axis),
1354
+ ],
1355
+ axis=1,
1356
+ )
1357
+ expected.columns = ["mean", "max", "sum"]
1358
+ expected = expected.T if axis in {0, "index"} else expected
1359
+
1360
+ result = float_frame.agg(["mean", "max", "sum"], axis=axis)
1361
+ tm.assert_frame_equal(result, expected)
1362
+
1363
+ # dict input with scalars
1364
+ func = {name1: "mean", name2: "sum"}
1365
+ result = float_frame.agg(func, axis=axis)
1366
+ expected = Series(
1367
+ [
1368
+ float_frame.loc(other_axis)[name1].mean(),
1369
+ float_frame.loc(other_axis)[name2].sum(),
1370
+ ],
1371
+ index=[name1, name2],
1372
+ )
1373
+ tm.assert_series_equal(result, expected)
1374
+
1375
+ # dict input with lists
1376
+ func = {name1: ["mean"], name2: ["sum"]}
1377
+ result = float_frame.agg(func, axis=axis)
1378
+ expected = DataFrame(
1379
+ {
1380
+ name1: Series([float_frame.loc(other_axis)[name1].mean()], index=["mean"]),
1381
+ name2: Series([float_frame.loc(other_axis)[name2].sum()], index=["sum"]),
1382
+ }
1383
+ )
1384
+ expected = expected.T if axis in {1, "columns"} else expected
1385
+ tm.assert_frame_equal(result, expected)
1386
+
1387
+ # dict input with lists with multiple
1388
+ func = {name1: ["mean", "sum"], name2: ["sum", "max"]}
1389
+ result = float_frame.agg(func, axis=axis)
1390
+ expected = pd.concat(
1391
+ {
1392
+ name1: Series(
1393
+ [
1394
+ float_frame.loc(other_axis)[name1].mean(),
1395
+ float_frame.loc(other_axis)[name1].sum(),
1396
+ ],
1397
+ index=["mean", "sum"],
1398
+ ),
1399
+ name2: Series(
1400
+ [
1401
+ float_frame.loc(other_axis)[name2].sum(),
1402
+ float_frame.loc(other_axis)[name2].max(),
1403
+ ],
1404
+ index=["sum", "max"],
1405
+ ),
1406
+ },
1407
+ axis=1,
1408
+ )
1409
+ expected = expected.T if axis in {1, "columns"} else expected
1410
+ tm.assert_frame_equal(result, expected)
1411
+
1412
+
1413
+ def test_named_agg_reduce_axis1_raises(float_frame):
1414
+ name1, name2 = float_frame.axes[0].unique()[:2].sort_values()
1415
+ msg = "Named aggregation is not supported when axis=1."
1416
+ for axis in [1, "columns"]:
1417
+ with pytest.raises(NotImplementedError, match=msg):
1418
+ float_frame.agg(row1=(name1, "sum"), row2=(name2, "max"), axis=axis)
1419
+
1420
+
1421
+ def test_nuiscance_columns():
1422
+ # GH 15015
1423
+ df = DataFrame(
1424
+ {
1425
+ "A": [1, 2, 3],
1426
+ "B": [1.0, 2.0, 3.0],
1427
+ "C": ["foo", "bar", "baz"],
1428
+ "D": date_range("20130101", periods=3),
1429
+ }
1430
+ )
1431
+
1432
+ result = df.agg("min")
1433
+ expected = Series([1, 1.0, "bar", Timestamp("20130101")], index=df.columns)
1434
+ tm.assert_series_equal(result, expected)
1435
+
1436
+ result = df.agg(["min"])
1437
+ expected = DataFrame(
1438
+ [[1, 1.0, "bar", Timestamp("20130101")]],
1439
+ index=["min"],
1440
+ columns=df.columns,
1441
+ )
1442
+ tm.assert_frame_equal(result, expected)
1443
+
1444
+ msg = "does not support operation"
1445
+ with pytest.raises(TypeError, match=msg):
1446
+ df.agg("sum")
1447
+
1448
+ result = df[["A", "B", "C"]].agg("sum")
1449
+ expected = Series([6, 6.0, "foobarbaz"], index=["A", "B", "C"])
1450
+ tm.assert_series_equal(result, expected)
1451
+
1452
+ msg = "does not support operation"
1453
+ with pytest.raises(TypeError, match=msg):
1454
+ df.agg(["sum"])
1455
+
1456
+
1457
+ @pytest.mark.parametrize("how", ["agg", "apply"])
1458
+ def test_non_callable_aggregates(how):
1459
+ # GH 16405
1460
+ # 'size' is a property of frame/series
1461
+ # validate that this is working
1462
+ # GH 39116 - expand to apply
1463
+ df = DataFrame(
1464
+ {"A": [None, 2, 3], "B": [1.0, np.nan, 3.0], "C": ["foo", None, "bar"]}
1465
+ )
1466
+
1467
+ # Function aggregate
1468
+ result = getattr(df, how)({"A": "count"})
1469
+ expected = Series({"A": 2})
1470
+
1471
+ tm.assert_series_equal(result, expected)
1472
+
1473
+ # Non-function aggregate
1474
+ result = getattr(df, how)({"A": "size"})
1475
+ expected = Series({"A": 3})
1476
+
1477
+ tm.assert_series_equal(result, expected)
1478
+
1479
+ # Mix function and non-function aggs
1480
+ result1 = getattr(df, how)(["count", "size"])
1481
+ result2 = getattr(df, how)(
1482
+ {"A": ["count", "size"], "B": ["count", "size"], "C": ["count", "size"]}
1483
+ )
1484
+ expected = DataFrame(
1485
+ {
1486
+ "A": {"count": 2, "size": 3},
1487
+ "B": {"count": 2, "size": 3},
1488
+ "C": {"count": 2, "size": 3},
1489
+ }
1490
+ )
1491
+
1492
+ tm.assert_frame_equal(result1, result2, check_like=True)
1493
+ tm.assert_frame_equal(result2, expected, check_like=True)
1494
+
1495
+ # Just functional string arg is same as calling df.arg()
1496
+ result = getattr(df, how)("count")
1497
+ expected = df.count()
1498
+
1499
+ tm.assert_series_equal(result, expected)
1500
+
1501
+
1502
+ @pytest.mark.parametrize("how", ["agg", "apply"])
1503
+ def test_size_as_str(how, axis):
1504
+ # GH 39934
1505
+ df = DataFrame(
1506
+ {"A": [None, 2, 3], "B": [1.0, np.nan, 3.0], "C": ["foo", None, "bar"]}
1507
+ )
1508
+ # Just a string attribute arg same as calling df.arg
1509
+ # on the columns
1510
+ result = getattr(df, how)("size", axis=axis)
1511
+ if axis in (0, "index"):
1512
+ expected = Series(df.shape[0], index=df.columns)
1513
+ else:
1514
+ expected = Series(df.shape[1], index=df.index)
1515
+ tm.assert_series_equal(result, expected)
1516
+
1517
+
1518
+ def test_agg_listlike_result():
1519
+ # GH-29587 user defined function returning list-likes
1520
+ df = DataFrame({"A": [2, 2, 3], "B": [1.5, np.nan, 1.5], "C": ["foo", None, "bar"]})
1521
+
1522
+ def func(group_col):
1523
+ return list(group_col.dropna().unique())
1524
+
1525
+ result = df.agg(func)
1526
+ expected = Series([[2, 3], [1.5], ["foo", "bar"]], index=["A", "B", "C"])
1527
+ tm.assert_series_equal(result, expected)
1528
+
1529
+ result = df.agg([func])
1530
+ expected = expected.to_frame("func").T
1531
+ tm.assert_frame_equal(result, expected)
1532
+
1533
+
1534
+ @pytest.mark.parametrize("axis", [0, 1])
1535
+ @pytest.mark.parametrize(
1536
+ "args, kwargs",
1537
+ [
1538
+ ((1, 2, 3), {}),
1539
+ ((8, 7, 15), {}),
1540
+ ((1, 2), {}),
1541
+ ((1,), {"b": 2}),
1542
+ ((), {"a": 1, "b": 2}),
1543
+ ((), {"a": 2, "b": 1}),
1544
+ ((), {"a": 1, "b": 2, "c": 3}),
1545
+ ],
1546
+ )
1547
+ def test_agg_args_kwargs(axis, args, kwargs):
1548
+ def f(x, a, b, c=3):
1549
+ return x.sum() + (a + b) / c
1550
+
1551
+ df = DataFrame([[1, 2], [3, 4]])
1552
+
1553
+ if axis == 0:
1554
+ expected = Series([5.0, 7.0])
1555
+ else:
1556
+ expected = Series([4.0, 8.0])
1557
+
1558
+ result = df.agg(f, axis, *args, **kwargs)
1559
+
1560
+ tm.assert_series_equal(result, expected)
1561
+
1562
+
1563
+ @pytest.mark.parametrize("num_cols", [2, 3, 5])
1564
+ def test_frequency_is_original(num_cols, engine, request):
1565
+ # GH 22150
1566
+ if engine == "numba":
1567
+ mark = pytest.mark.xfail(reason="numba engine only supports numeric indices")
1568
+ request.node.add_marker(mark)
1569
+ index = pd.DatetimeIndex(["1950-06-30", "1952-10-24", "1953-05-29"])
1570
+ original = index.copy()
1571
+ df = DataFrame(1, index=index, columns=range(num_cols))
1572
+ df.apply(lambda x: x, engine=engine)
1573
+ assert index.freq == original.freq
1574
+
1575
+
1576
+ def test_apply_datetime_tz_issue(engine, request):
1577
+ # GH 29052
1578
+
1579
+ if engine == "numba":
1580
+ mark = pytest.mark.xfail(
1581
+ reason="numba engine doesn't support non-numeric indexes"
1582
+ )
1583
+ request.node.add_marker(mark)
1584
+
1585
+ timestamps = [
1586
+ Timestamp("2019-03-15 12:34:31.909000+0000", tz="UTC"),
1587
+ Timestamp("2019-03-15 12:34:34.359000+0000", tz="UTC"),
1588
+ Timestamp("2019-03-15 12:34:34.660000+0000", tz="UTC"),
1589
+ ]
1590
+ df = DataFrame(data=[0, 1, 2], index=timestamps)
1591
+ result = df.apply(lambda x: x.name, axis=1, engine=engine)
1592
+ expected = Series(index=timestamps, data=timestamps)
1593
+
1594
+ tm.assert_series_equal(result, expected)
1595
+
1596
+
1597
+ @pytest.mark.parametrize("df", [DataFrame({"A": ["a", None], "B": ["c", "d"]})])
1598
+ @pytest.mark.parametrize("method", ["min", "max", "sum"])
1599
+ def test_mixed_column_raises(df, method, using_infer_string):
1600
+ # GH 16832
1601
+ if method == "sum":
1602
+ msg = r'can only concatenate str \(not "int"\) to str|does not support'
1603
+ else:
1604
+ msg = "not supported between instances of 'str' and 'float'"
1605
+ if not using_infer_string:
1606
+ with pytest.raises(TypeError, match=msg):
1607
+ getattr(df, method)()
1608
+ else:
1609
+ getattr(df, method)()
1610
+
1611
+
1612
+ @pytest.mark.parametrize("col", [1, 1.0, True, "a", np.nan])
1613
+ def test_apply_dtype(col):
1614
+ # GH 31466
1615
+ df = DataFrame([[1.0, col]], columns=["a", "b"])
1616
+ result = df.apply(lambda x: x.dtype)
1617
+ expected = df.dtypes
1618
+
1619
+ tm.assert_series_equal(result, expected)
1620
+
1621
+
1622
+ def test_apply_mutating():
1623
+ # GH#35462 case where applied func pins a new BlockManager to a row
1624
+ df = DataFrame({"a": range(10), "b": range(10, 20)})
1625
+ df_orig = df.copy()
1626
+
1627
+ def func(row):
1628
+ mgr = row._mgr
1629
+ row.loc["a"] += 1
1630
+ assert row._mgr is not mgr
1631
+ return row
1632
+
1633
+ expected = df.copy()
1634
+ expected["a"] += 1
1635
+
1636
+ result = df.apply(func, axis=1)
1637
+
1638
+ tm.assert_frame_equal(result, expected)
1639
+ tm.assert_frame_equal(df, df_orig)
1640
+
1641
+
1642
+ def test_apply_empty_list_reduce():
1643
+ # GH#35683 get columns correct
1644
+ df = DataFrame([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]], columns=["a", "b"])
1645
+
1646
+ result = df.apply(lambda x: [], result_type="reduce")
1647
+ expected = Series({"a": [], "b": []}, dtype=object)
1648
+ tm.assert_series_equal(result, expected)
1649
+
1650
+
1651
+ def test_apply_no_suffix_index(engine, request):
1652
+ # GH36189
1653
+ if engine == "numba":
1654
+ mark = pytest.mark.xfail(
1655
+ reason="numba engine doesn't support list-likes/dict-like callables"
1656
+ )
1657
+ request.node.add_marker(mark)
1658
+ pdf = DataFrame([[4, 9]] * 3, columns=["A", "B"])
1659
+ result = pdf.apply(["sum", lambda x: x.sum(), lambda x: x.sum()], engine=engine)
1660
+ expected = DataFrame(
1661
+ {"A": [12, 12, 12], "B": [27, 27, 27]}, index=["sum", "<lambda>", "<lambda>"]
1662
+ )
1663
+
1664
+ tm.assert_frame_equal(result, expected)
1665
+
1666
+
1667
+ def test_apply_raw_returns_string(engine):
1668
+ # https://github.com/pandas-dev/pandas/issues/35940
1669
+ if engine == "numba":
1670
+ pytest.skip("No object dtype support in numba")
1671
+ df = DataFrame({"A": ["aa", "bbb"]})
1672
+ result = df.apply(lambda x: x[0], engine=engine, axis=1, raw=True)
1673
+ expected = Series(["aa", "bbb"])
1674
+ tm.assert_series_equal(result, expected)
1675
+
1676
+
1677
+ def test_aggregation_func_column_order():
1678
+ # GH40420: the result of .agg should have an index that is sorted
1679
+ # according to the arguments provided to agg.
1680
+ df = DataFrame(
1681
+ [
1682
+ (1, 0, 0),
1683
+ (2, 0, 0),
1684
+ (3, 0, 0),
1685
+ (4, 5, 4),
1686
+ (5, 6, 6),
1687
+ (6, 7, 7),
1688
+ ],
1689
+ columns=("att1", "att2", "att3"),
1690
+ )
1691
+
1692
+ def sum_div2(s):
1693
+ return s.sum() / 2
1694
+
1695
+ aggs = ["sum", sum_div2, "count", "min"]
1696
+ result = df.agg(aggs)
1697
+ expected = DataFrame(
1698
+ {
1699
+ "att1": [21.0, 10.5, 6.0, 1.0],
1700
+ "att2": [18.0, 9.0, 6.0, 0.0],
1701
+ "att3": [17.0, 8.5, 6.0, 0.0],
1702
+ },
1703
+ index=["sum", "sum_div2", "count", "min"],
1704
+ )
1705
+ tm.assert_frame_equal(result, expected)
1706
+
1707
+
1708
+ def test_apply_getitem_axis_1(engine, request):
1709
+ # GH 13427
1710
+ if engine == "numba":
1711
+ mark = pytest.mark.xfail(
1712
+ reason="numba engine not supporting duplicate index values"
1713
+ )
1714
+ request.node.add_marker(mark)
1715
+ df = DataFrame({"a": [0, 1, 2], "b": [1, 2, 3]})
1716
+ result = df[["a", "a"]].apply(
1717
+ lambda x: x.iloc[0] + x.iloc[1], axis=1, engine=engine
1718
+ )
1719
+ expected = Series([0, 2, 4])
1720
+ tm.assert_series_equal(result, expected)
1721
+
1722
+
1723
+ def test_nuisance_depr_passes_through_warnings():
1724
+ # GH 43740
1725
+ # DataFrame.agg with list-likes may emit warnings for both individual
1726
+ # args and for entire columns, but we only want to emit once. We
1727
+ # catch and suppress the warnings for individual args, but need to make
1728
+ # sure if some other warnings were raised, they get passed through to
1729
+ # the user.
1730
+
1731
+ def expected_warning(x):
1732
+ warnings.warn("Hello, World!")
1733
+ return x.sum()
1734
+
1735
+ df = DataFrame({"a": [1, 2, 3]})
1736
+ with tm.assert_produces_warning(UserWarning, match="Hello, World!"):
1737
+ df.agg([expected_warning])
1738
+
1739
+
1740
+ def test_apply_type():
1741
+ # GH 46719
1742
+ df = DataFrame(
1743
+ {"col1": [3, "string", float], "col2": [0.25, datetime(2020, 1, 1), np.nan]},
1744
+ index=["a", "b", "c"],
1745
+ )
1746
+
1747
+ # axis=0
1748
+ result = df.apply(type, axis=0)
1749
+ expected = Series({"col1": Series, "col2": Series})
1750
+ tm.assert_series_equal(result, expected)
1751
+
1752
+ # axis=1
1753
+ result = df.apply(type, axis=1)
1754
+ expected = Series({"a": Series, "b": Series, "c": Series})
1755
+ tm.assert_series_equal(result, expected)
1756
+
1757
+
1758
+ def test_apply_on_empty_dataframe(engine):
1759
+ # GH 39111
1760
+ df = DataFrame({"a": [1, 2], "b": [3, 0]})
1761
+ result = df.head(0).apply(lambda x: max(x["a"], x["b"]), axis=1, engine=engine)
1762
+ expected = Series([], dtype=np.float64)
1763
+ tm.assert_series_equal(result, expected)
1764
+
1765
+
1766
+ def test_apply_return_list():
1767
+ df = DataFrame({"a": [1, 2], "b": [2, 3]})
1768
+ result = df.apply(lambda x: [x.values])
1769
+ expected = DataFrame({"a": [[1, 2]], "b": [[2, 3]]})
1770
+ tm.assert_frame_equal(result, expected)
1771
+
1772
+
1773
+ @pytest.mark.parametrize(
1774
+ "test, constant",
1775
+ [
1776
+ ({"a": [1, 2, 3], "b": [1, 1, 1]}, {"a": [1, 2, 3], "b": [1]}),
1777
+ ({"a": [2, 2, 2], "b": [1, 1, 1]}, {"a": [2], "b": [1]}),
1778
+ ],
1779
+ )
1780
+ def test_unique_agg_type_is_series(test, constant):
1781
+ # GH#22558
1782
+ df1 = DataFrame(test)
1783
+ expected = Series(data=constant, index=["a", "b"], dtype="object")
1784
+ aggregation = {"a": "unique", "b": "unique"}
1785
+
1786
+ result = df1.agg(aggregation)
1787
+
1788
+ tm.assert_series_equal(result, expected)
1789
+
1790
+
1791
+ def test_any_apply_keyword_non_zero_axis_regression():
1792
+ # https://github.com/pandas-dev/pandas/issues/48656
1793
+ df = DataFrame({"A": [1, 2, 0], "B": [0, 2, 0], "C": [0, 0, 0]})
1794
+ expected = Series([True, True, False])
1795
+ tm.assert_series_equal(df.any(axis=1), expected)
1796
+
1797
+ result = df.apply("any", axis=1)
1798
+ tm.assert_series_equal(result, expected)
1799
+
1800
+ result = df.apply("any", 1)
1801
+ tm.assert_series_equal(result, expected)
1802
+
1803
+
1804
+ def test_agg_mapping_func_deprecated():
1805
+ # GH 53325
1806
+ df = DataFrame({"x": [1, 2, 3]})
1807
+
1808
+ def foo1(x, a=1, c=0):
1809
+ return x + a + c
1810
+
1811
+ def foo2(x, b=2, c=0):
1812
+ return x + b + c
1813
+
1814
+ # single func already takes the vectorized path
1815
+ result = df.agg(foo1, 0, 3, c=4)
1816
+ expected = df + 7
1817
+ tm.assert_frame_equal(result, expected)
1818
+
1819
+ result = df.agg([foo1, foo2], 0, 3, c=4)
1820
+ expected = DataFrame(
1821
+ [[8, 8], [9, 9], [10, 10]], columns=[["x", "x"], ["foo1", "foo2"]]
1822
+ )
1823
+ tm.assert_frame_equal(result, expected)
1824
+
1825
+ # TODO: the result below is wrong, should be fixed (GH53325)
1826
+ result = df.agg({"x": foo1}, 0, 3, c=4)
1827
+ expected = DataFrame([2, 3, 4], columns=["x"])
1828
+ tm.assert_frame_equal(result, expected)
1829
+
1830
+
1831
+ def test_agg_std():
1832
+ df = DataFrame(np.arange(6).reshape(3, 2), columns=["A", "B"])
1833
+
1834
+ result = df.agg(np.std, ddof=1)
1835
+ expected = Series({"A": 2.0, "B": 2.0}, dtype=float)
1836
+ tm.assert_series_equal(result, expected)
1837
+
1838
+ result = df.agg([np.std], ddof=1)
1839
+ expected = DataFrame({"A": 2.0, "B": 2.0}, index=["std"])
1840
+ tm.assert_frame_equal(result, expected)
1841
+
1842
+
1843
+ def test_agg_np_size():
1844
+ # GH#42203, GH#48328
1845
+ df = DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=["A", "B", "C"])
1846
+
1847
+ result = df.agg({"A": [np.size]})
1848
+ expected = DataFrame({"A": [3]}, index=["size"])
1849
+ tm.assert_frame_equal(result, expected)
1850
+
1851
+ result = df.agg({"A": np.size})
1852
+ expected = Series({"A": 3})
1853
+ tm.assert_series_equal(result, expected)
1854
+
1855
+ result = df.agg({"A": [np.mean, np.size]})
1856
+ expected = DataFrame({"A": [4.0, 3.0]}, index=["mean", "size"])
1857
+ tm.assert_frame_equal(result, expected)
1858
+
1859
+
1860
+ def test_agg_dist_like_and_nonunique_columns():
1861
+ # GH#51099
1862
+ df = DataFrame(
1863
+ {"A": [None, 2, 3], "B": [1.0, np.nan, 3.0], "C": ["foo", None, "bar"]}
1864
+ )
1865
+ df.columns = ["A", "A", "C"]
1866
+
1867
+ result = df.agg({"A": "count"})
1868
+ expected = df["A"].count()
1869
+ tm.assert_series_equal(result, expected)
1870
+
1871
+
1872
+ @pytest.mark.parametrize("engine_name", ["unknown", 25])
1873
+ def test_wrong_engine(engine_name):
1874
+ with pytest.raises(ValueError, match="Unknown engine "):
1875
+ DataFrame().apply(lambda x: x, engine=engine_name)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_apply_relabeling.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ import pandas as pd
4
+ import pandas._testing as tm
5
+
6
+
7
+ def test_agg_relabel():
8
+ # GH 26513
9
+ df = pd.DataFrame({"A": [1, 2, 1, 2], "B": [1, 2, 3, 4], "C": [3, 4, 5, 6]})
10
+
11
+ # simplest case with one column, one func
12
+ result = df.agg(foo=("B", "sum"))
13
+ expected = pd.DataFrame({"B": [10]}, index=pd.Index(["foo"]))
14
+ tm.assert_frame_equal(result, expected)
15
+
16
+ # test on same column with different methods
17
+ result = df.agg(foo=("B", "sum"), bar=("B", "min"))
18
+ expected = pd.DataFrame({"B": [10, 1]}, index=pd.Index(["foo", "bar"]))
19
+
20
+ tm.assert_frame_equal(result, expected)
21
+
22
+
23
+ def test_agg_relabel_multi_columns_multi_methods():
24
+ # GH 26513, test on multiple columns with multiple methods
25
+ df = pd.DataFrame({"A": [1, 2, 1, 2], "B": [1, 2, 3, 4], "C": [3, 4, 5, 6]})
26
+ result = df.agg(
27
+ foo=("A", "sum"),
28
+ bar=("B", "mean"),
29
+ cat=("A", "min"),
30
+ dat=("B", "max"),
31
+ f=("A", "max"),
32
+ g=("C", "min"),
33
+ )
34
+ expected = pd.DataFrame(
35
+ {
36
+ "A": [6.0, np.nan, 1.0, np.nan, 2.0, np.nan],
37
+ "B": [np.nan, 2.5, np.nan, 4.0, np.nan, np.nan],
38
+ "C": [np.nan, np.nan, np.nan, np.nan, np.nan, 3.0],
39
+ },
40
+ index=pd.Index(["foo", "bar", "cat", "dat", "f", "g"]),
41
+ )
42
+ tm.assert_frame_equal(result, expected)
43
+
44
+
45
+ def test_agg_relabel_partial_functions():
46
+ # GH 26513, test on partial, functools or more complex cases
47
+ df = pd.DataFrame({"A": [1, 2, 1, 2], "B": [1, 2, 3, 4], "C": [3, 4, 5, 6]})
48
+ result = df.agg(foo=("A", np.mean), bar=("A", "mean"), cat=("A", min))
49
+ expected = pd.DataFrame(
50
+ {"A": [1.5, 1.5, 1.0]}, index=pd.Index(["foo", "bar", "cat"])
51
+ )
52
+ tm.assert_frame_equal(result, expected)
53
+
54
+ result = df.agg(
55
+ foo=("A", min),
56
+ bar=("B", np.min),
57
+ cat=("B", max),
58
+ dat=("C", "min"),
59
+ f=("B", np.sum),
60
+ kk=("B", lambda x: min(x)),
61
+ )
62
+ expected = pd.DataFrame(
63
+ {
64
+ "A": [1.0, np.nan, np.nan, np.nan, np.nan, np.nan],
65
+ "B": [np.nan, 1.0, 4.0, np.nan, 10.0, 1.0],
66
+ "C": [np.nan, np.nan, np.nan, 3.0, np.nan, np.nan],
67
+ },
68
+ index=pd.Index(["foo", "bar", "cat", "dat", "f", "kk"]),
69
+ )
70
+ tm.assert_frame_equal(result, expected)
71
+
72
+
73
+ def test_agg_namedtuple():
74
+ # GH 26513
75
+ df = pd.DataFrame({"A": [0, 1], "B": [1, 2]})
76
+ result = df.agg(
77
+ foo=pd.NamedAgg("B", "sum"),
78
+ bar=pd.NamedAgg("B", "min"),
79
+ cat=pd.NamedAgg(column="B", aggfunc="count"),
80
+ fft=pd.NamedAgg("B", aggfunc="max"),
81
+ )
82
+
83
+ expected = pd.DataFrame(
84
+ {"B": [3, 1, 2, 2]}, index=pd.Index(["foo", "bar", "cat", "fft"])
85
+ )
86
+ tm.assert_frame_equal(result, expected)
87
+
88
+ result = df.agg(
89
+ foo=pd.NamedAgg("A", "min"),
90
+ bar=pd.NamedAgg(column="B", aggfunc="max"),
91
+ cat=pd.NamedAgg(column="A", aggfunc="max"),
92
+ )
93
+ expected = pd.DataFrame(
94
+ {"A": [0.0, np.nan, 1.0], "B": [np.nan, 2.0, np.nan]},
95
+ index=pd.Index(["foo", "bar", "cat"]),
96
+ )
97
+ tm.assert_frame_equal(result, expected)
98
+
99
+
100
+ def test_reconstruct_func():
101
+ # GH 28472, test to ensure reconstruct_func isn't moved;
102
+ # This method is used by other libraries (e.g. dask)
103
+ result = pd.core.apply.reconstruct_func("min")
104
+ expected = (False, "min", None, None)
105
+ tm.assert_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_frame_transform.py ADDED
@@ -0,0 +1,264 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas import (
5
+ DataFrame,
6
+ MultiIndex,
7
+ Series,
8
+ )
9
+ import pandas._testing as tm
10
+ from pandas.tests.apply.common import frame_transform_kernels
11
+ from pandas.tests.frame.common import zip_frames
12
+
13
+
14
+ def unpack_obj(obj, klass, axis):
15
+ """
16
+ Helper to ensure we have the right type of object for a test parametrized
17
+ over frame_or_series.
18
+ """
19
+ if klass is not DataFrame:
20
+ obj = obj["A"]
21
+ if axis != 0:
22
+ pytest.skip(f"Test is only for DataFrame with axis={axis}")
23
+ return obj
24
+
25
+
26
+ def test_transform_ufunc(axis, float_frame, frame_or_series):
27
+ # GH 35964
28
+ obj = unpack_obj(float_frame, frame_or_series, axis)
29
+
30
+ with np.errstate(all="ignore"):
31
+ f_sqrt = np.sqrt(obj)
32
+
33
+ # ufunc
34
+ result = obj.transform(np.sqrt, axis=axis)
35
+ expected = f_sqrt
36
+ tm.assert_equal(result, expected)
37
+
38
+
39
+ @pytest.mark.parametrize(
40
+ "ops, names",
41
+ [
42
+ ([np.sqrt], ["sqrt"]),
43
+ ([np.abs, np.sqrt], ["absolute", "sqrt"]),
44
+ (np.array([np.sqrt]), ["sqrt"]),
45
+ (np.array([np.abs, np.sqrt]), ["absolute", "sqrt"]),
46
+ ],
47
+ )
48
+ def test_transform_listlike(axis, float_frame, ops, names):
49
+ # GH 35964
50
+ other_axis = 1 if axis in {0, "index"} else 0
51
+ with np.errstate(all="ignore"):
52
+ expected = zip_frames([op(float_frame) for op in ops], axis=other_axis)
53
+ if axis in {0, "index"}:
54
+ expected.columns = MultiIndex.from_product([float_frame.columns, names])
55
+ else:
56
+ expected.index = MultiIndex.from_product([float_frame.index, names])
57
+ result = float_frame.transform(ops, axis=axis)
58
+ tm.assert_frame_equal(result, expected)
59
+
60
+
61
+ @pytest.mark.parametrize("ops", [[], np.array([])])
62
+ def test_transform_empty_listlike(float_frame, ops, frame_or_series):
63
+ obj = unpack_obj(float_frame, frame_or_series, 0)
64
+
65
+ with pytest.raises(ValueError, match="No transform functions were provided"):
66
+ obj.transform(ops)
67
+
68
+
69
+ def test_transform_listlike_func_with_args():
70
+ # GH 50624
71
+ df = DataFrame({"x": [1, 2, 3]})
72
+
73
+ def foo1(x, a=1, c=0):
74
+ return x + a + c
75
+
76
+ def foo2(x, b=2, c=0):
77
+ return x + b + c
78
+
79
+ msg = r"foo1\(\) got an unexpected keyword argument 'b'"
80
+ with pytest.raises(TypeError, match=msg):
81
+ df.transform([foo1, foo2], 0, 3, b=3, c=4)
82
+
83
+ result = df.transform([foo1, foo2], 0, 3, c=4)
84
+ expected = DataFrame(
85
+ [[8, 8], [9, 9], [10, 10]],
86
+ columns=MultiIndex.from_tuples([("x", "foo1"), ("x", "foo2")]),
87
+ )
88
+ tm.assert_frame_equal(result, expected)
89
+
90
+
91
+ @pytest.mark.parametrize("box", [dict, Series])
92
+ def test_transform_dictlike(axis, float_frame, box):
93
+ # GH 35964
94
+ if axis in (0, "index"):
95
+ e = float_frame.columns[0]
96
+ expected = float_frame[[e]].transform(np.abs)
97
+ else:
98
+ e = float_frame.index[0]
99
+ expected = float_frame.iloc[[0]].transform(np.abs)
100
+ result = float_frame.transform(box({e: np.abs}), axis=axis)
101
+ tm.assert_frame_equal(result, expected)
102
+
103
+
104
+ def test_transform_dictlike_mixed():
105
+ # GH 40018 - mix of lists and non-lists in values of a dictionary
106
+ df = DataFrame({"a": [1, 2], "b": [1, 4], "c": [1, 4]})
107
+ result = df.transform({"b": ["sqrt", "abs"], "c": "sqrt"})
108
+ expected = DataFrame(
109
+ [[1.0, 1, 1.0], [2.0, 4, 2.0]],
110
+ columns=MultiIndex([("b", "c"), ("sqrt", "abs")], [(0, 0, 1), (0, 1, 0)]),
111
+ )
112
+ tm.assert_frame_equal(result, expected)
113
+
114
+
115
+ @pytest.mark.parametrize(
116
+ "ops",
117
+ [
118
+ {},
119
+ {"A": []},
120
+ {"A": [], "B": "cumsum"},
121
+ {"A": "cumsum", "B": []},
122
+ {"A": [], "B": ["cumsum"]},
123
+ {"A": ["cumsum"], "B": []},
124
+ ],
125
+ )
126
+ def test_transform_empty_dictlike(float_frame, ops, frame_or_series):
127
+ obj = unpack_obj(float_frame, frame_or_series, 0)
128
+
129
+ with pytest.raises(ValueError, match="No transform functions were provided"):
130
+ obj.transform(ops)
131
+
132
+
133
+ @pytest.mark.parametrize("use_apply", [True, False])
134
+ def test_transform_udf(axis, float_frame, use_apply, frame_or_series):
135
+ # GH 35964
136
+ obj = unpack_obj(float_frame, frame_or_series, axis)
137
+
138
+ # transform uses UDF either via apply or passing the entire DataFrame
139
+ def func(x):
140
+ # transform is using apply iff x is not a DataFrame
141
+ if use_apply == isinstance(x, frame_or_series):
142
+ # Force transform to fallback
143
+ raise ValueError
144
+ return x + 1
145
+
146
+ result = obj.transform(func, axis=axis)
147
+ expected = obj + 1
148
+ tm.assert_equal(result, expected)
149
+
150
+
151
+ wont_fail = ["ffill", "bfill", "fillna", "pad", "backfill", "shift"]
152
+ frame_kernels_raise = [x for x in frame_transform_kernels if x not in wont_fail]
153
+
154
+
155
+ @pytest.mark.parametrize("op", [*frame_kernels_raise, lambda x: x + 1])
156
+ def test_transform_bad_dtype(op, frame_or_series, request):
157
+ # GH 35964
158
+ if op == "ngroup":
159
+ request.applymarker(
160
+ pytest.mark.xfail(raises=ValueError, reason="ngroup not valid for NDFrame")
161
+ )
162
+
163
+ obj = DataFrame({"A": 3 * [object]}) # DataFrame that will fail on most transforms
164
+ obj = tm.get_obj(obj, frame_or_series)
165
+ error = TypeError
166
+ msg = "|".join(
167
+ [
168
+ "not supported between instances of 'type' and 'type'",
169
+ "unsupported operand type",
170
+ ]
171
+ )
172
+
173
+ with pytest.raises(error, match=msg):
174
+ obj.transform(op)
175
+ with pytest.raises(error, match=msg):
176
+ obj.transform([op])
177
+ with pytest.raises(error, match=msg):
178
+ obj.transform({"A": op})
179
+ with pytest.raises(error, match=msg):
180
+ obj.transform({"A": [op]})
181
+
182
+
183
+ @pytest.mark.parametrize("op", frame_kernels_raise)
184
+ def test_transform_failure_typeerror(request, op):
185
+ # GH 35964
186
+
187
+ if op == "ngroup":
188
+ request.applymarker(
189
+ pytest.mark.xfail(raises=ValueError, reason="ngroup not valid for NDFrame")
190
+ )
191
+
192
+ # Using object makes most transform kernels fail
193
+ df = DataFrame({"A": 3 * [object], "B": [1, 2, 3]})
194
+ error = TypeError
195
+ msg = "|".join(
196
+ [
197
+ "not supported between instances of 'type' and 'type'",
198
+ "unsupported operand type",
199
+ ]
200
+ )
201
+
202
+ with pytest.raises(error, match=msg):
203
+ df.transform([op])
204
+
205
+ with pytest.raises(error, match=msg):
206
+ df.transform({"A": op, "B": op})
207
+
208
+ with pytest.raises(error, match=msg):
209
+ df.transform({"A": [op], "B": [op]})
210
+
211
+ with pytest.raises(error, match=msg):
212
+ df.transform({"A": [op, "shift"], "B": [op]})
213
+
214
+
215
+ def test_transform_failure_valueerror():
216
+ # GH 40211
217
+ def op(x):
218
+ if np.sum(np.sum(x)) < 10:
219
+ raise ValueError
220
+ return x
221
+
222
+ df = DataFrame({"A": [1, 2, 3], "B": [400, 500, 600]})
223
+ msg = "Transform function failed"
224
+
225
+ with pytest.raises(ValueError, match=msg):
226
+ df.transform([op])
227
+
228
+ with pytest.raises(ValueError, match=msg):
229
+ df.transform({"A": op, "B": op})
230
+
231
+ with pytest.raises(ValueError, match=msg):
232
+ df.transform({"A": [op], "B": [op]})
233
+
234
+ with pytest.raises(ValueError, match=msg):
235
+ df.transform({"A": [op, "shift"], "B": [op]})
236
+
237
+
238
+ @pytest.mark.parametrize("use_apply", [True, False])
239
+ def test_transform_passes_args(use_apply, frame_or_series):
240
+ # GH 35964
241
+ # transform uses UDF either via apply or passing the entire DataFrame
242
+ expected_args = [1, 2]
243
+ expected_kwargs = {"c": 3}
244
+
245
+ def f(x, a, b, c):
246
+ # transform is using apply iff x is not a DataFrame
247
+ if use_apply == isinstance(x, frame_or_series):
248
+ # Force transform to fallback
249
+ raise ValueError
250
+ assert [a, b] == expected_args
251
+ assert c == expected_kwargs["c"]
252
+ return x
253
+
254
+ frame_or_series([1]).transform(f, 0, *expected_args, **expected_kwargs)
255
+
256
+
257
+ def test_transform_empty_dataframe():
258
+ # https://github.com/pandas-dev/pandas/issues/39636
259
+ df = DataFrame([], columns=["col1", "col2"])
260
+ result = df.transform(lambda x: x + 10)
261
+ tm.assert_frame_equal(result, df)
262
+
263
+ result = df["col1"].transform(lambda x: x + 10)
264
+ tm.assert_series_equal(result, df["col1"])
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_invalid_arg.py ADDED
@@ -0,0 +1,375 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Tests specifically aimed at detecting bad arguments.
2
+ # This file is organized by reason for exception.
3
+ # 1. always invalid argument values
4
+ # 2. missing column(s)
5
+ # 3. incompatible ops/dtype/args/kwargs
6
+ # 4. invalid result shape/type
7
+ # If your test does not fit into one of these categories, add to this list.
8
+
9
+ from itertools import chain
10
+ import re
11
+
12
+ import numpy as np
13
+ import pytest
14
+
15
+ from pandas.errors import SpecificationError
16
+
17
+ from pandas import (
18
+ DataFrame,
19
+ Series,
20
+ date_range,
21
+ )
22
+ import pandas._testing as tm
23
+
24
+
25
+ @pytest.mark.parametrize("result_type", ["foo", 1])
26
+ def test_result_type_error(result_type):
27
+ # allowed result_type
28
+ df = DataFrame(
29
+ np.tile(np.arange(3, dtype="int64"), 6).reshape(6, -1) + 1,
30
+ columns=["A", "B", "C"],
31
+ )
32
+
33
+ msg = (
34
+ "invalid value for result_type, must be one of "
35
+ "{None, 'reduce', 'broadcast', 'expand'}"
36
+ )
37
+ with pytest.raises(ValueError, match=msg):
38
+ df.apply(lambda x: [1, 2, 3], axis=1, result_type=result_type)
39
+
40
+
41
+ def test_apply_invalid_axis_value():
42
+ df = DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=["a", "a", "c"])
43
+ msg = "No axis named 2 for object type DataFrame"
44
+ with pytest.raises(ValueError, match=msg):
45
+ df.apply(lambda x: x, 2)
46
+
47
+
48
+ def test_agg_raises():
49
+ # GH 26513
50
+ df = DataFrame({"A": [0, 1], "B": [1, 2]})
51
+ msg = "Must provide"
52
+
53
+ with pytest.raises(TypeError, match=msg):
54
+ df.agg()
55
+
56
+
57
+ def test_map_with_invalid_na_action_raises():
58
+ # https://github.com/pandas-dev/pandas/issues/32815
59
+ s = Series([1, 2, 3])
60
+ msg = "na_action must either be 'ignore' or None"
61
+ with pytest.raises(ValueError, match=msg):
62
+ s.map(lambda x: x, na_action="____")
63
+
64
+
65
+ @pytest.mark.parametrize("input_na_action", ["____", True])
66
+ def test_map_arg_is_dict_with_invalid_na_action_raises(input_na_action):
67
+ # https://github.com/pandas-dev/pandas/issues/46588
68
+ s = Series([1, 2, 3])
69
+ msg = f"na_action must either be 'ignore' or None, {input_na_action} was passed"
70
+ with pytest.raises(ValueError, match=msg):
71
+ s.map({1: 2}, na_action=input_na_action)
72
+
73
+
74
+ @pytest.mark.parametrize("method", ["apply", "agg", "transform"])
75
+ @pytest.mark.parametrize("func", [{"A": {"B": "sum"}}, {"A": {"B": ["sum"]}}])
76
+ def test_nested_renamer(frame_or_series, method, func):
77
+ # GH 35964
78
+ obj = frame_or_series({"A": [1]})
79
+ match = "nested renamer is not supported"
80
+ with pytest.raises(SpecificationError, match=match):
81
+ getattr(obj, method)(func)
82
+
83
+
84
+ @pytest.mark.parametrize(
85
+ "renamer",
86
+ [{"foo": ["min", "max"]}, {"foo": ["min", "max"], "bar": ["sum", "mean"]}],
87
+ )
88
+ def test_series_nested_renamer(renamer):
89
+ s = Series(range(6), dtype="int64", name="series")
90
+ msg = "nested renamer is not supported"
91
+ with pytest.raises(SpecificationError, match=msg):
92
+ s.agg(renamer)
93
+
94
+
95
+ def test_apply_dict_depr():
96
+ tsdf = DataFrame(
97
+ np.random.default_rng(2).standard_normal((10, 3)),
98
+ columns=["A", "B", "C"],
99
+ index=date_range("1/1/2000", periods=10),
100
+ )
101
+ msg = "nested renamer is not supported"
102
+ with pytest.raises(SpecificationError, match=msg):
103
+ tsdf.A.agg({"foo": ["sum", "mean"]})
104
+
105
+
106
+ @pytest.mark.parametrize("method", ["agg", "transform"])
107
+ def test_dict_nested_renaming_depr(method):
108
+ df = DataFrame({"A": range(5), "B": 5})
109
+
110
+ # nested renaming
111
+ msg = r"nested renamer is not supported"
112
+ with pytest.raises(SpecificationError, match=msg):
113
+ getattr(df, method)({"A": {"foo": "min"}, "B": {"bar": "max"}})
114
+
115
+
116
+ @pytest.mark.parametrize("method", ["apply", "agg", "transform"])
117
+ @pytest.mark.parametrize("func", [{"B": "sum"}, {"B": ["sum"]}])
118
+ def test_missing_column(method, func):
119
+ # GH 40004
120
+ obj = DataFrame({"A": [1]})
121
+ msg = r"Label\(s\) \['B'\] do not exist"
122
+ with pytest.raises(KeyError, match=msg):
123
+ getattr(obj, method)(func)
124
+
125
+
126
+ def test_transform_mixed_column_name_dtypes():
127
+ # GH39025
128
+ df = DataFrame({"a": ["1"]})
129
+ msg = r"Label\(s\) \[1, 'b'\] do not exist"
130
+ with pytest.raises(KeyError, match=msg):
131
+ df.transform({"a": int, 1: str, "b": int})
132
+
133
+
134
+ @pytest.mark.parametrize(
135
+ "how, args", [("pct_change", ()), ("nsmallest", (1, ["a", "b"])), ("tail", 1)]
136
+ )
137
+ def test_apply_str_axis_1_raises(how, args):
138
+ # GH 39211 - some ops don't support axis=1
139
+ df = DataFrame({"a": [1, 2], "b": [3, 4]})
140
+ msg = f"Operation {how} does not support axis=1"
141
+ with pytest.raises(ValueError, match=msg):
142
+ df.apply(how, axis=1, args=args)
143
+
144
+
145
+ def test_transform_axis_1_raises():
146
+ # GH 35964
147
+ msg = "No axis named 1 for object type Series"
148
+ with pytest.raises(ValueError, match=msg):
149
+ Series([1]).transform("sum", axis=1)
150
+
151
+
152
+ def test_apply_modify_traceback():
153
+ data = DataFrame(
154
+ {
155
+ "A": [
156
+ "foo",
157
+ "foo",
158
+ "foo",
159
+ "foo",
160
+ "bar",
161
+ "bar",
162
+ "bar",
163
+ "bar",
164
+ "foo",
165
+ "foo",
166
+ "foo",
167
+ ],
168
+ "B": [
169
+ "one",
170
+ "one",
171
+ "one",
172
+ "two",
173
+ "one",
174
+ "one",
175
+ "one",
176
+ "two",
177
+ "two",
178
+ "two",
179
+ "one",
180
+ ],
181
+ "C": [
182
+ "dull",
183
+ "dull",
184
+ "shiny",
185
+ "dull",
186
+ "dull",
187
+ "shiny",
188
+ "shiny",
189
+ "dull",
190
+ "shiny",
191
+ "shiny",
192
+ "shiny",
193
+ ],
194
+ "D": np.random.default_rng(2).standard_normal(11),
195
+ "E": np.random.default_rng(2).standard_normal(11),
196
+ "F": np.random.default_rng(2).standard_normal(11),
197
+ }
198
+ )
199
+
200
+ data.loc[4, "C"] = np.nan
201
+
202
+ def transform(row):
203
+ if row["C"].startswith("shin") and row["A"] == "foo":
204
+ row["D"] = 7
205
+ return row
206
+
207
+ msg = "'float' object has no attribute 'startswith'"
208
+ with pytest.raises(AttributeError, match=msg):
209
+ data.apply(transform, axis=1)
210
+
211
+
212
+ @pytest.mark.parametrize(
213
+ "df, func, expected",
214
+ tm.get_cython_table_params(
215
+ DataFrame([["a", "b"], ["b", "a"]]), [["cumprod", TypeError]]
216
+ ),
217
+ )
218
+ def test_agg_cython_table_raises_frame(df, func, expected, axis, using_infer_string):
219
+ # GH 21224
220
+ if using_infer_string:
221
+ expected = (expected, NotImplementedError)
222
+
223
+ msg = (
224
+ "can't multiply sequence by non-int of type 'str'"
225
+ "|cannot perform cumprod with type str" # NotImplementedError python backend
226
+ "|operation 'cumprod' not supported for dtype 'str'" # TypeError pyarrow
227
+ )
228
+ warn = None if isinstance(func, str) else FutureWarning
229
+ with pytest.raises(expected, match=msg):
230
+ with tm.assert_produces_warning(warn, match="using DataFrame.cumprod"):
231
+ df.agg(func, axis=axis)
232
+
233
+
234
+ @pytest.mark.parametrize(
235
+ "series, func, expected",
236
+ chain(
237
+ tm.get_cython_table_params(
238
+ Series("a b c".split()),
239
+ [
240
+ ("mean", TypeError), # mean raises TypeError
241
+ ("prod", TypeError),
242
+ ("std", TypeError),
243
+ ("var", TypeError),
244
+ ("median", TypeError),
245
+ ("cumprod", TypeError),
246
+ ],
247
+ )
248
+ ),
249
+ )
250
+ def test_agg_cython_table_raises_series(series, func, expected, using_infer_string):
251
+ # GH21224
252
+ msg = r"[Cc]ould not convert|can't multiply sequence by non-int of type"
253
+ if func == "median" or func is np.nanmedian or func is np.median:
254
+ msg = r"Cannot convert \['a' 'b' 'c'\] to numeric"
255
+
256
+ if using_infer_string and func == "cumprod":
257
+ expected = (expected, NotImplementedError)
258
+
259
+ msg = (
260
+ msg + "|does not support|has no kernel|Cannot perform|cannot perform|operation"
261
+ )
262
+ warn = None if isinstance(func, str) else FutureWarning
263
+
264
+ with pytest.raises(expected, match=msg):
265
+ # e.g. Series('a b'.split()).cumprod() will raise
266
+ with tm.assert_produces_warning(warn, match="is currently using Series.*"):
267
+ series.agg(func)
268
+
269
+
270
+ def test_agg_none_to_type():
271
+ # GH 40543
272
+ df = DataFrame({"a": [None]})
273
+ msg = re.escape("int() argument must be a string")
274
+ with pytest.raises(TypeError, match=msg):
275
+ df.agg({"a": lambda x: int(x.iloc[0])})
276
+
277
+
278
+ def test_transform_none_to_type():
279
+ # GH#34377
280
+ df = DataFrame({"a": [None]})
281
+ msg = "argument must be a"
282
+ with pytest.raises(TypeError, match=msg):
283
+ df.transform({"a": lambda x: int(x.iloc[0])})
284
+
285
+
286
+ @pytest.mark.parametrize(
287
+ "func",
288
+ [
289
+ lambda x: np.array([1, 2]).reshape(-1, 2),
290
+ lambda x: [1, 2],
291
+ lambda x: Series([1, 2]),
292
+ ],
293
+ )
294
+ def test_apply_broadcast_error(func):
295
+ df = DataFrame(
296
+ np.tile(np.arange(3, dtype="int64"), 6).reshape(6, -1) + 1,
297
+ columns=["A", "B", "C"],
298
+ )
299
+
300
+ # > 1 ndim
301
+ msg = "too many dims to broadcast|cannot broadcast result"
302
+ with pytest.raises(ValueError, match=msg):
303
+ df.apply(func, axis=1, result_type="broadcast")
304
+
305
+
306
+ def test_transform_and_agg_err_agg(axis, float_frame):
307
+ # cannot both transform and agg
308
+ msg = "cannot combine transform and aggregation operations"
309
+ with pytest.raises(ValueError, match=msg):
310
+ with np.errstate(all="ignore"):
311
+ float_frame.agg(["max", "sqrt"], axis=axis)
312
+
313
+
314
+ @pytest.mark.filterwarnings("ignore::FutureWarning") # GH53325
315
+ @pytest.mark.parametrize(
316
+ "func, msg",
317
+ [
318
+ (["sqrt", "max"], "cannot combine transform and aggregation"),
319
+ (
320
+ {"foo": np.sqrt, "bar": "sum"},
321
+ "cannot perform both aggregation and transformation",
322
+ ),
323
+ ],
324
+ )
325
+ def test_transform_and_agg_err_series(string_series, func, msg):
326
+ # we are trying to transform with an aggregator
327
+ with pytest.raises(ValueError, match=msg):
328
+ with np.errstate(all="ignore"):
329
+ string_series.agg(func)
330
+
331
+
332
+ @pytest.mark.parametrize("func", [["max", "min"], ["max", "sqrt"]])
333
+ def test_transform_wont_agg_frame(axis, float_frame, func):
334
+ # GH 35964
335
+ # cannot both transform and agg
336
+ msg = "Function did not transform"
337
+ with pytest.raises(ValueError, match=msg):
338
+ float_frame.transform(func, axis=axis)
339
+
340
+
341
+ @pytest.mark.parametrize("func", [["min", "max"], ["sqrt", "max"]])
342
+ def test_transform_wont_agg_series(string_series, func):
343
+ # GH 35964
344
+ # we are trying to transform with an aggregator
345
+ msg = "Function did not transform"
346
+
347
+ with pytest.raises(ValueError, match=msg):
348
+ string_series.transform(func)
349
+
350
+
351
+ @pytest.mark.parametrize(
352
+ "op_wrapper", [lambda x: x, lambda x: [x], lambda x: {"A": x}, lambda x: {"A": [x]}]
353
+ )
354
+ def test_transform_reducer_raises(all_reductions, frame_or_series, op_wrapper):
355
+ # GH 35964
356
+ op = op_wrapper(all_reductions)
357
+
358
+ obj = DataFrame({"A": [1, 2, 3]})
359
+ obj = tm.get_obj(obj, frame_or_series)
360
+
361
+ msg = "Function did not transform"
362
+ with pytest.raises(ValueError, match=msg):
363
+ obj.transform(op)
364
+
365
+
366
+ def test_transform_missing_labels_raises():
367
+ # GH 58474
368
+ df = DataFrame({"foo": [2, 4, 6], "bar": [1, 2, 3]}, index=["A", "B", "C"])
369
+ msg = r"Label\(s\) \['A', 'B'\] do not exist"
370
+ with pytest.raises(KeyError, match=msg):
371
+ df.transform({"A": lambda x: x + 2, "B": lambda x: x * 2}, axis=0)
372
+
373
+ msg = r"Label\(s\) \['bar', 'foo'\] do not exist"
374
+ with pytest.raises(KeyError, match=msg):
375
+ df.transform({"foo": lambda x: x + 2, "bar": lambda x: x * 2}, axis=1)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_numba.py ADDED
@@ -0,0 +1,129 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas.compat import is_platform_arm
5
+ import pandas.util._test_decorators as td
6
+
7
+ import pandas as pd
8
+ from pandas import (
9
+ DataFrame,
10
+ Index,
11
+ )
12
+ import pandas._testing as tm
13
+ from pandas.util.version import Version
14
+
15
+ pytestmark = [td.skip_if_no("numba"), pytest.mark.single_cpu, pytest.mark.skipif()]
16
+
17
+ numba = pytest.importorskip("numba")
18
+ pytestmark.append(
19
+ pytest.mark.skipif(
20
+ Version(numba.__version__) == Version("0.61") and is_platform_arm(),
21
+ reason=f"Segfaults on ARM platforms with numba {numba.__version__}",
22
+ )
23
+ )
24
+
25
+
26
+ @pytest.fixture(params=[0, 1])
27
+ def apply_axis(request):
28
+ return request.param
29
+
30
+
31
+ def test_numba_vs_python_noop(float_frame, apply_axis):
32
+ func = lambda x: x
33
+ result = float_frame.apply(func, engine="numba", axis=apply_axis)
34
+ expected = float_frame.apply(func, engine="python", axis=apply_axis)
35
+ tm.assert_frame_equal(result, expected)
36
+
37
+
38
+ def test_numba_vs_python_string_index():
39
+ # GH#56189
40
+ df = DataFrame(
41
+ 1,
42
+ index=Index(["a", "b"], dtype=pd.StringDtype(na_value=np.nan)),
43
+ columns=Index(["x", "y"], dtype=pd.StringDtype(na_value=np.nan)),
44
+ )
45
+ func = lambda x: x
46
+ result = df.apply(func, engine="numba", axis=0)
47
+ expected = df.apply(func, engine="python", axis=0)
48
+ tm.assert_frame_equal(
49
+ result, expected, check_column_type=False, check_index_type=False
50
+ )
51
+
52
+
53
+ def test_numba_vs_python_indexing():
54
+ frame = DataFrame(
55
+ {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7.0, 8.0, 9.0]},
56
+ index=Index(["A", "B", "C"]),
57
+ )
58
+ row_func = lambda x: x["c"]
59
+ result = frame.apply(row_func, engine="numba", axis=1)
60
+ expected = frame.apply(row_func, engine="python", axis=1)
61
+ tm.assert_series_equal(result, expected)
62
+
63
+ col_func = lambda x: x["A"]
64
+ result = frame.apply(col_func, engine="numba", axis=0)
65
+ expected = frame.apply(col_func, engine="python", axis=0)
66
+ tm.assert_series_equal(result, expected)
67
+
68
+
69
+ @pytest.mark.parametrize(
70
+ "reduction",
71
+ [lambda x: x.mean(), lambda x: x.min(), lambda x: x.max(), lambda x: x.sum()],
72
+ )
73
+ def test_numba_vs_python_reductions(reduction, apply_axis):
74
+ df = DataFrame(np.ones((4, 4), dtype=np.float64))
75
+ result = df.apply(reduction, engine="numba", axis=apply_axis)
76
+ expected = df.apply(reduction, engine="python", axis=apply_axis)
77
+ tm.assert_series_equal(result, expected)
78
+
79
+
80
+ @pytest.mark.parametrize("colnames", [[1, 2, 3], [1.0, 2.0, 3.0]])
81
+ def test_numba_numeric_colnames(colnames):
82
+ # Check that numeric column names lower properly and can be indexed on
83
+ df = DataFrame(
84
+ np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.int64), columns=colnames
85
+ )
86
+ first_col = colnames[0]
87
+ f = lambda x: x[first_col] # Get the first column
88
+ result = df.apply(f, engine="numba", axis=1)
89
+ expected = df.apply(f, engine="python", axis=1)
90
+ tm.assert_series_equal(result, expected)
91
+
92
+
93
+ def test_numba_parallel_unsupported(float_frame):
94
+ f = lambda x: x
95
+ with pytest.raises(
96
+ NotImplementedError,
97
+ match="Parallel apply is not supported when raw=False and engine='numba'",
98
+ ):
99
+ float_frame.apply(f, engine="numba", engine_kwargs={"parallel": True})
100
+
101
+
102
+ def test_numba_nonunique_unsupported(apply_axis):
103
+ f = lambda x: x
104
+ df = DataFrame({"a": [1, 2]}, index=Index(["a", "a"]))
105
+ with pytest.raises(
106
+ NotImplementedError,
107
+ match="The index/columns must be unique when raw=False and engine='numba'",
108
+ ):
109
+ df.apply(f, engine="numba", axis=apply_axis)
110
+
111
+
112
+ def test_numba_unsupported_dtypes(apply_axis):
113
+ pytest.importorskip("pyarrow")
114
+ f = lambda x: x
115
+ df = DataFrame({"a": [1, 2], "b": ["a", "b"], "c": [4, 5]})
116
+ df["c"] = df["c"].astype("double[pyarrow]")
117
+
118
+ with pytest.raises(
119
+ ValueError,
120
+ match="Column b must have a numeric dtype. Found 'object|str' instead",
121
+ ):
122
+ df.apply(f, engine="numba", axis=apply_axis)
123
+
124
+ with pytest.raises(
125
+ ValueError,
126
+ match="Column c is backed by an extension array, "
127
+ "which is not supported by the numba engine.",
128
+ ):
129
+ df["c"].to_frame().apply(f, engine="numba", axis=apply_axis)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_apply.py ADDED
@@ -0,0 +1,669 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ from pandas import (
6
+ DataFrame,
7
+ Index,
8
+ MultiIndex,
9
+ Series,
10
+ concat,
11
+ date_range,
12
+ timedelta_range,
13
+ )
14
+ import pandas._testing as tm
15
+ from pandas.tests.apply.common import series_transform_kernels
16
+
17
+
18
+ @pytest.fixture(params=[False, "compat"])
19
+ def by_row(request):
20
+ return request.param
21
+
22
+
23
+ def test_series_map_box_timedelta(by_row):
24
+ # GH#11349
25
+ ser = Series(timedelta_range("1 day 1 s", periods=3, freq="h"))
26
+
27
+ def f(x):
28
+ return x.total_seconds() if by_row else x.dt.total_seconds()
29
+
30
+ result = ser.apply(f, by_row=by_row)
31
+
32
+ expected = ser.map(lambda x: x.total_seconds())
33
+ tm.assert_series_equal(result, expected)
34
+
35
+ expected = Series([86401.0, 90001.0, 93601.0])
36
+ tm.assert_series_equal(result, expected)
37
+
38
+
39
+ def test_apply(datetime_series, by_row):
40
+ result = datetime_series.apply(np.sqrt, by_row=by_row)
41
+ with np.errstate(all="ignore"):
42
+ expected = np.sqrt(datetime_series)
43
+ tm.assert_series_equal(result, expected)
44
+
45
+ # element-wise apply (ufunc)
46
+ result = datetime_series.apply(np.exp, by_row=by_row)
47
+ expected = np.exp(datetime_series)
48
+ tm.assert_series_equal(result, expected)
49
+
50
+ # empty series
51
+ s = Series(dtype=object, name="foo", index=Index([], name="bar"))
52
+ rs = s.apply(lambda x: x, by_row=by_row)
53
+ tm.assert_series_equal(s, rs)
54
+
55
+ # check all metadata (GH 9322)
56
+ assert s is not rs
57
+ assert s.index is rs.index
58
+ assert s.dtype == rs.dtype
59
+ assert s.name == rs.name
60
+
61
+ # index but no data
62
+ s = Series(index=[1, 2, 3], dtype=np.float64)
63
+ rs = s.apply(lambda x: x, by_row=by_row)
64
+ tm.assert_series_equal(s, rs)
65
+
66
+
67
+ def test_apply_map_same_length_inference_bug():
68
+ s = Series([1, 2])
69
+
70
+ def f(x):
71
+ return (x, x + 1)
72
+
73
+ result = s.apply(f, by_row="compat")
74
+ expected = s.map(f)
75
+ tm.assert_series_equal(result, expected)
76
+
77
+
78
+ def test_apply_args():
79
+ s = Series(["foo,bar"])
80
+
81
+ result = s.apply(str.split, args=(",",))
82
+ assert result[0] == ["foo", "bar"]
83
+ assert isinstance(result[0], list)
84
+
85
+
86
+ @pytest.mark.parametrize(
87
+ "args, kwargs, increment",
88
+ [((), {}, 0), ((), {"a": 1}, 1), ((2, 3), {}, 32), ((1,), {"c": 2}, 201)],
89
+ )
90
+ def test_agg_args(args, kwargs, increment):
91
+ # GH 43357
92
+ def f(x, a=0, b=0, c=0):
93
+ return x + a + 10 * b + 100 * c
94
+
95
+ s = Series([1, 2])
96
+ result = s.agg(f, 0, *args, **kwargs)
97
+ expected = s + increment
98
+ tm.assert_series_equal(result, expected)
99
+
100
+
101
+ def test_agg_mapping_func_deprecated():
102
+ # GH 53325
103
+ s = Series([1, 2, 3])
104
+
105
+ def foo1(x, a=1, c=0):
106
+ return x + a + c
107
+
108
+ def foo2(x, b=2, c=0):
109
+ return x + b + c
110
+
111
+ s.agg(foo1, 0, 3, c=4)
112
+ s.agg([foo1, foo2], 0, 3, c=4)
113
+ s.agg({"a": foo1, "b": foo2}, 0, 3, c=4)
114
+
115
+
116
+ def test_series_apply_map_box_timestamps(by_row):
117
+ # GH#2689, GH#2627
118
+ ser = Series(date_range("1/1/2000", periods=10))
119
+
120
+ def func(x):
121
+ return (x.hour, x.day, x.month)
122
+
123
+ if not by_row:
124
+ msg = "Series' object has no attribute 'hour'"
125
+ with pytest.raises(AttributeError, match=msg):
126
+ ser.apply(func, by_row=by_row)
127
+ return
128
+
129
+ result = ser.apply(func, by_row=by_row)
130
+ expected = ser.map(func)
131
+ tm.assert_series_equal(result, expected)
132
+
133
+
134
+ def test_apply_box_dt64():
135
+ # ufunc will not be boxed. Same test cases as the test_map_box
136
+ vals = [pd.Timestamp("2011-01-01"), pd.Timestamp("2011-01-02")]
137
+ ser = Series(vals, dtype="M8[ns]")
138
+ assert ser.dtype == "datetime64[ns]"
139
+ # boxed value must be Timestamp instance
140
+ res = ser.apply(lambda x: f"{type(x).__name__}_{x.day}_{x.tz}", by_row="compat")
141
+ exp = Series(["Timestamp_1_None", "Timestamp_2_None"])
142
+ tm.assert_series_equal(res, exp)
143
+
144
+
145
+ def test_apply_box_dt64tz():
146
+ vals = [
147
+ pd.Timestamp("2011-01-01", tz="US/Eastern"),
148
+ pd.Timestamp("2011-01-02", tz="US/Eastern"),
149
+ ]
150
+ ser = Series(vals, dtype="M8[ns, US/Eastern]")
151
+ assert ser.dtype == "datetime64[ns, US/Eastern]"
152
+ res = ser.apply(lambda x: f"{type(x).__name__}_{x.day}_{x.tz}", by_row="compat")
153
+ exp = Series(["Timestamp_1_US/Eastern", "Timestamp_2_US/Eastern"])
154
+ tm.assert_series_equal(res, exp)
155
+
156
+
157
+ def test_apply_box_td64():
158
+ # timedelta
159
+ vals = [pd.Timedelta("1 days"), pd.Timedelta("2 days")]
160
+ ser = Series(vals)
161
+ assert ser.dtype == "timedelta64[us]"
162
+ res = ser.apply(lambda x: f"{type(x).__name__}_{x.days}", by_row="compat")
163
+ exp = Series(["Timedelta_1", "Timedelta_2"])
164
+ tm.assert_series_equal(res, exp)
165
+
166
+
167
+ def test_apply_box_period():
168
+ # period
169
+ vals = [pd.Period("2011-01-01", freq="M"), pd.Period("2011-01-02", freq="M")]
170
+ ser = Series(vals)
171
+ assert ser.dtype == "Period[M]"
172
+ res = ser.apply(lambda x: f"{type(x).__name__}_{x.freqstr}", by_row="compat")
173
+ exp = Series(["Period_M", "Period_M"])
174
+ tm.assert_series_equal(res, exp)
175
+
176
+
177
+ def test_apply_datetimetz(by_row):
178
+ values = date_range("2011-01-01", "2011-01-02", freq="h").tz_localize("Asia/Tokyo")
179
+ s = Series(values, name="XX")
180
+
181
+ result = s.apply(lambda x: x + pd.offsets.Day(), by_row=by_row)
182
+ exp_values = date_range("2011-01-02", "2011-01-03", freq="h").tz_localize(
183
+ "Asia/Tokyo"
184
+ )
185
+ exp = Series(exp_values, name="XX")
186
+ tm.assert_series_equal(result, exp)
187
+
188
+ result = s.apply(lambda x: x.hour if by_row else x.dt.hour, by_row=by_row)
189
+ exp = Series([*list(range(24)), 0], name="XX", dtype="int64" if by_row else "int32")
190
+ tm.assert_series_equal(result, exp)
191
+
192
+ # not vectorized
193
+ def f(x):
194
+ return str(x.tz) if by_row else str(x.dt.tz)
195
+
196
+ result = s.apply(f, by_row=by_row)
197
+ if by_row:
198
+ exp = Series(["Asia/Tokyo"] * 25, name="XX")
199
+ tm.assert_series_equal(result, exp)
200
+ else:
201
+ assert result == "Asia/Tokyo"
202
+
203
+
204
+ def test_apply_categorical(by_row, using_infer_string):
205
+ values = pd.Categorical(list("ABBABCD"), categories=list("DCBA"), ordered=True)
206
+ ser = Series(values, name="XX", index=list("abcdefg"))
207
+
208
+ if not by_row:
209
+ msg = "Series' object has no attribute 'lower"
210
+ with pytest.raises(AttributeError, match=msg):
211
+ ser.apply(lambda x: x.lower(), by_row=by_row)
212
+ assert ser.apply(lambda x: "A", by_row=by_row) == "A"
213
+ return
214
+
215
+ result = ser.apply(lambda x: x.lower(), by_row=by_row)
216
+
217
+ # should be categorical dtype when the number of categories are
218
+ # the same
219
+ values = pd.Categorical(list("abbabcd"), categories=list("dcba"), ordered=True)
220
+ exp = Series(values, name="XX", index=list("abcdefg"))
221
+ tm.assert_series_equal(result, exp)
222
+ tm.assert_categorical_equal(result.values, exp.values)
223
+
224
+ result = ser.apply(lambda x: "A")
225
+ exp = Series(["A"] * 7, name="XX", index=list("abcdefg"))
226
+ tm.assert_series_equal(result, exp)
227
+ assert result.dtype == object if not using_infer_string else "str"
228
+
229
+
230
+ @pytest.mark.parametrize("series", [["1-1", "1-1", np.nan], ["1-1", "1-2", np.nan]])
231
+ def test_apply_categorical_with_nan_values(series, by_row):
232
+ # GH 20714 bug fixed in: GH 24275
233
+ s = Series(series, dtype="category")
234
+ if not by_row:
235
+ msg = "'Series' object has no attribute 'split'"
236
+ with pytest.raises(AttributeError, match=msg):
237
+ s.apply(lambda x: x.split("-")[0], by_row=by_row)
238
+ return
239
+ # NaN for cat dtype fixed in (GH 59966)
240
+ result = s.apply(lambda x: x.split("-")[0] if pd.notna(x) else False, by_row=by_row)
241
+ result = result.astype(object)
242
+ expected = Series(["1", "1", False], dtype="category")
243
+ expected = expected.astype(object)
244
+ tm.assert_series_equal(result, expected)
245
+
246
+
247
+ def test_apply_empty_integer_series_with_datetime_index(by_row):
248
+ # GH 21245
249
+ s = Series([], index=date_range(start="2018-01-01", periods=0), dtype=int)
250
+ result = s.apply(lambda x: x, by_row=by_row)
251
+ tm.assert_series_equal(result, s)
252
+
253
+
254
+ def test_apply_dataframe_iloc():
255
+ uintDF = DataFrame(np.uint64([1, 2, 3, 4, 5]), columns=["Numbers"])
256
+ indexDF = DataFrame([2, 3, 2, 1, 2], columns=["Indices"])
257
+
258
+ def retrieve(targetRow, targetDF):
259
+ val = targetDF["Numbers"].iloc[targetRow]
260
+ return val
261
+
262
+ result = indexDF["Indices"].apply(retrieve, args=(uintDF,))
263
+ expected = Series([3, 4, 3, 2, 3], name="Indices", dtype="uint64")
264
+ tm.assert_series_equal(result, expected)
265
+
266
+
267
+ def test_transform(string_series, by_row):
268
+ # transforming functions
269
+
270
+ with np.errstate(all="ignore"):
271
+ f_sqrt = np.sqrt(string_series)
272
+ f_abs = np.abs(string_series)
273
+
274
+ # ufunc
275
+ result = string_series.apply(np.sqrt, by_row=by_row)
276
+ expected = f_sqrt.copy()
277
+ tm.assert_series_equal(result, expected)
278
+
279
+ # list-like
280
+ result = string_series.apply([np.sqrt], by_row=by_row)
281
+ expected = f_sqrt.to_frame().copy()
282
+ expected.columns = ["sqrt"]
283
+ tm.assert_frame_equal(result, expected)
284
+
285
+ result = string_series.apply(["sqrt"], by_row=by_row)
286
+ tm.assert_frame_equal(result, expected)
287
+
288
+ # multiple items in list
289
+ # these are in the order as if we are applying both functions per
290
+ # series and then concatting
291
+ expected = concat([f_sqrt, f_abs], axis=1)
292
+ expected.columns = ["sqrt", "absolute"]
293
+ result = string_series.apply([np.sqrt, np.abs], by_row=by_row)
294
+ tm.assert_frame_equal(result, expected)
295
+
296
+ # dict, provide renaming
297
+ expected = concat([f_sqrt, f_abs], axis=1)
298
+ expected.columns = ["foo", "bar"]
299
+ expected = expected.unstack().rename("series")
300
+
301
+ result = string_series.apply({"foo": np.sqrt, "bar": np.abs}, by_row=by_row)
302
+ tm.assert_series_equal(result.reindex_like(expected), expected)
303
+
304
+
305
+ @pytest.mark.parametrize("op", series_transform_kernels)
306
+ def test_transform_partial_failure(op, request):
307
+ # GH 35964
308
+ if op in ("ffill", "bfill", "shift"):
309
+ request.applymarker(
310
+ pytest.mark.xfail(reason=f"{op} is successful on any dtype")
311
+ )
312
+
313
+ # Using object makes most transform kernels fail
314
+ ser = Series(3 * [object])
315
+
316
+ if op in ("fillna", "ngroup"):
317
+ error = ValueError
318
+ msg = "Transform function failed"
319
+ else:
320
+ error = TypeError
321
+ msg = "|".join(
322
+ [
323
+ "not supported between instances of 'type' and 'type'",
324
+ "unsupported operand type",
325
+ ]
326
+ )
327
+
328
+ with pytest.raises(error, match=msg):
329
+ ser.transform([op, "shift"])
330
+
331
+ with pytest.raises(error, match=msg):
332
+ ser.transform({"A": op, "B": "shift"})
333
+
334
+ with pytest.raises(error, match=msg):
335
+ ser.transform({"A": [op], "B": ["shift"]})
336
+
337
+ with pytest.raises(error, match=msg):
338
+ ser.transform({"A": [op, "shift"], "B": [op]})
339
+
340
+
341
+ def test_transform_partial_failure_valueerror():
342
+ # GH 40211
343
+ def noop(x):
344
+ return x
345
+
346
+ def raising_op(_):
347
+ raise ValueError
348
+
349
+ ser = Series(3 * [object])
350
+ msg = "Transform function failed"
351
+
352
+ with pytest.raises(ValueError, match=msg):
353
+ ser.transform([noop, raising_op])
354
+
355
+ with pytest.raises(ValueError, match=msg):
356
+ ser.transform({"A": raising_op, "B": noop})
357
+
358
+ with pytest.raises(ValueError, match=msg):
359
+ ser.transform({"A": [raising_op], "B": [noop]})
360
+
361
+ with pytest.raises(ValueError, match=msg):
362
+ ser.transform({"A": [noop, raising_op], "B": [noop]})
363
+
364
+
365
+ def test_demo():
366
+ # demonstration tests
367
+ s = Series(range(6), dtype="int64", name="series")
368
+
369
+ result = s.agg(["min", "max"])
370
+ expected = Series([0, 5], index=["min", "max"], name="series")
371
+ tm.assert_series_equal(result, expected)
372
+
373
+ result = s.agg({"foo": "min"})
374
+ expected = Series([0], index=["foo"], name="series")
375
+ tm.assert_series_equal(result, expected)
376
+
377
+
378
+ @pytest.mark.parametrize("func", [str, lambda x: str(x)])
379
+ def test_apply_map_evaluate_lambdas_the_same(string_series, func, by_row, engine):
380
+ # test that we are evaluating row-by-row first if by_row="compat"
381
+ # else vectorized evaluation
382
+ result = string_series.apply(func, by_row=by_row)
383
+
384
+ if by_row:
385
+ expected = string_series.map(func, engine=engine)
386
+ tm.assert_series_equal(result, expected)
387
+ else:
388
+ assert result == str(string_series)
389
+
390
+
391
+ def test_agg_evaluate_lambdas(string_series):
392
+ # GH53325
393
+ result = string_series.agg(lambda x: type(x))
394
+ assert result is Series
395
+
396
+ result = string_series.agg(type)
397
+ assert result is Series
398
+
399
+
400
+ @pytest.mark.parametrize("op_name", ["agg", "apply"])
401
+ def test_with_nested_series(datetime_series, op_name):
402
+ # GH 2316 & GH52123
403
+ # .agg with a reducer and a transform, what to do
404
+ result = getattr(datetime_series, op_name)(
405
+ lambda x: Series([x, x**2], index=["x", "x^2"])
406
+ )
407
+ if op_name == "apply":
408
+ expected = DataFrame({"x": datetime_series, "x^2": datetime_series**2})
409
+ tm.assert_frame_equal(result, expected)
410
+ else:
411
+ expected = Series([datetime_series, datetime_series**2], index=["x", "x^2"])
412
+ tm.assert_series_equal(result, expected)
413
+
414
+
415
+ def test_replicate_describe(string_series):
416
+ # this also tests a result set that is all scalars
417
+ expected = string_series.describe()
418
+ result = string_series.apply(
419
+ {
420
+ "count": "count",
421
+ "mean": "mean",
422
+ "std": "std",
423
+ "min": "min",
424
+ "25%": lambda x: x.quantile(0.25),
425
+ "50%": "median",
426
+ "75%": lambda x: x.quantile(0.75),
427
+ "max": "max",
428
+ },
429
+ )
430
+ tm.assert_series_equal(result, expected)
431
+
432
+
433
+ def test_reduce(string_series):
434
+ # reductions with named functions
435
+ result = string_series.agg(["sum", "mean"])
436
+ expected = Series(
437
+ [string_series.sum(), string_series.mean()],
438
+ ["sum", "mean"],
439
+ name=string_series.name,
440
+ )
441
+ tm.assert_series_equal(result, expected)
442
+
443
+
444
+ @pytest.mark.parametrize(
445
+ "how, kwds",
446
+ [("agg", {}), ("apply", {"by_row": "compat"}), ("apply", {"by_row": False})],
447
+ )
448
+ def test_non_callable_aggregates(how, kwds):
449
+ # test agg using non-callable series attributes
450
+ # GH 39116 - expand to apply
451
+ s = Series([1, 2, None])
452
+
453
+ # Calling agg w/ just a string arg same as calling s.arg
454
+ result = getattr(s, how)("size", **kwds)
455
+ expected = s.size
456
+ assert result == expected
457
+
458
+ # test when mixed w/ callable reducers
459
+ result = getattr(s, how)(["size", "count", "mean"], **kwds)
460
+ expected = Series({"size": 3.0, "count": 2.0, "mean": 1.5})
461
+ tm.assert_series_equal(result, expected)
462
+
463
+ result = getattr(s, how)({"size": "size", "count": "count", "mean": "mean"}, **kwds)
464
+ tm.assert_series_equal(result, expected)
465
+
466
+
467
+ def test_series_apply_no_suffix_index(by_row):
468
+ # GH36189
469
+ s = Series([4] * 3)
470
+ result = s.apply(["sum", lambda x: x.sum(), lambda x: x.sum()], by_row=by_row)
471
+ expected = Series([12, 12, 12], index=["sum", "<lambda>", "<lambda>"])
472
+
473
+ tm.assert_series_equal(result, expected)
474
+
475
+
476
+ @pytest.mark.parametrize(
477
+ "dti,exp",
478
+ [
479
+ (
480
+ Series([1, 2], index=pd.DatetimeIndex([0, 31536000000])),
481
+ DataFrame(np.repeat([[1, 2]], 2, axis=0), dtype="int64"),
482
+ ),
483
+ (
484
+ Series(
485
+ np.arange(10, dtype=np.float64),
486
+ index=date_range("2020-01-01", periods=10),
487
+ name="ts",
488
+ ),
489
+ DataFrame(np.repeat([[1, 2]], 10, axis=0), dtype="int64"),
490
+ ),
491
+ ],
492
+ )
493
+ @pytest.mark.parametrize("aware", [True, False])
494
+ def test_apply_series_on_date_time_index_aware_series(dti, exp, aware):
495
+ # GH 25959
496
+ # Calling apply on a localized time series should not cause an error
497
+ if aware:
498
+ index = dti.tz_localize("UTC").index
499
+ else:
500
+ index = dti.index
501
+ result = Series(index).apply(lambda x: Series([1, 2]))
502
+ tm.assert_frame_equal(result, exp)
503
+
504
+
505
+ @pytest.mark.parametrize(
506
+ "by_row, expected", [("compat", Series(np.ones(10), dtype="int64")), (False, 1)]
507
+ )
508
+ def test_apply_scalar_on_date_time_index_aware_series(by_row, expected):
509
+ # GH 25959
510
+ # Calling apply on a localized time series should not cause an error
511
+ series = Series(
512
+ np.arange(10, dtype=np.float64),
513
+ index=date_range("2020-01-01", periods=10, tz="UTC"),
514
+ )
515
+ result = Series(series.index).apply(lambda x: 1, by_row=by_row)
516
+ tm.assert_equal(result, expected)
517
+
518
+
519
+ def test_apply_to_timedelta(by_row):
520
+ list_of_valid_strings = ["00:00:01", "00:00:02"]
521
+ a = pd.to_timedelta(list_of_valid_strings)
522
+ b = Series(list_of_valid_strings).apply(pd.to_timedelta, by_row=by_row)
523
+ tm.assert_series_equal(Series(a), b)
524
+
525
+ list_of_strings = ["00:00:01", np.nan, pd.NaT, pd.NaT]
526
+
527
+ a = pd.to_timedelta(list_of_strings)
528
+ ser = Series(list_of_strings)
529
+ b = ser.apply(pd.to_timedelta, by_row=by_row)
530
+ tm.assert_series_equal(Series(a), b)
531
+
532
+
533
+ @pytest.mark.parametrize(
534
+ "ops, names",
535
+ [
536
+ ([np.sum], ["sum"]),
537
+ ([np.sum, np.mean], ["sum", "mean"]),
538
+ (np.array([np.sum]), ["sum"]),
539
+ (np.array([np.sum, np.mean]), ["sum", "mean"]),
540
+ ],
541
+ )
542
+ @pytest.mark.parametrize(
543
+ "how, kwargs",
544
+ [["agg", {}], ["apply", {"by_row": "compat"}], ["apply", {"by_row": False}]],
545
+ )
546
+ def test_apply_listlike_reducer(string_series, ops, names, how, kwargs):
547
+ # GH 39140
548
+ expected = Series(
549
+ {name: op(string_series) for name, op in zip(names, ops, strict=True)}
550
+ )
551
+ expected.name = "series"
552
+ result = getattr(string_series, how)(ops, **kwargs)
553
+ tm.assert_series_equal(result, expected)
554
+
555
+
556
+ @pytest.mark.parametrize(
557
+ "ops",
558
+ [
559
+ {"A": np.sum},
560
+ {"A": np.sum, "B": np.mean},
561
+ Series({"A": np.sum}),
562
+ Series({"A": np.sum, "B": np.mean}),
563
+ ],
564
+ )
565
+ @pytest.mark.parametrize(
566
+ "how, kwargs",
567
+ [["agg", {}], ["apply", {"by_row": "compat"}], ["apply", {"by_row": False}]],
568
+ )
569
+ def test_apply_dictlike_reducer(string_series, ops, how, kwargs, by_row):
570
+ # GH 39140
571
+ expected = Series({name: op(string_series) for name, op in ops.items()})
572
+ expected.name = string_series.name
573
+ result = getattr(string_series, how)(ops, **kwargs)
574
+ tm.assert_series_equal(result, expected)
575
+
576
+
577
+ @pytest.mark.parametrize(
578
+ "ops, names",
579
+ [
580
+ ([np.sqrt], ["sqrt"]),
581
+ ([np.abs, np.sqrt], ["absolute", "sqrt"]),
582
+ (np.array([np.sqrt]), ["sqrt"]),
583
+ (np.array([np.abs, np.sqrt]), ["absolute", "sqrt"]),
584
+ ],
585
+ )
586
+ def test_apply_listlike_transformer(string_series, ops, names, by_row):
587
+ # GH 39140
588
+ with np.errstate(all="ignore"):
589
+ expected = concat([op(string_series) for op in ops], axis=1)
590
+ expected.columns = names
591
+ result = string_series.apply(ops, by_row=by_row)
592
+ tm.assert_frame_equal(result, expected)
593
+
594
+
595
+ @pytest.mark.parametrize(
596
+ "ops, expected",
597
+ [
598
+ ([lambda x: x], DataFrame({"<lambda>": [1, 2, 3]})),
599
+ ([lambda x: x.sum()], Series([6], index=["<lambda>"])),
600
+ ],
601
+ )
602
+ def test_apply_listlike_lambda(ops, expected, by_row):
603
+ # GH53400
604
+ ser = Series([1, 2, 3])
605
+ result = ser.apply(ops, by_row=by_row)
606
+ tm.assert_equal(result, expected)
607
+
608
+
609
+ @pytest.mark.parametrize(
610
+ "ops",
611
+ [
612
+ {"A": np.sqrt},
613
+ {"A": np.sqrt, "B": np.exp},
614
+ Series({"A": np.sqrt}),
615
+ Series({"A": np.sqrt, "B": np.exp}),
616
+ ],
617
+ )
618
+ def test_apply_dictlike_transformer(string_series, ops, by_row):
619
+ # GH 39140
620
+ with np.errstate(all="ignore"):
621
+ expected = concat({name: op(string_series) for name, op in ops.items()})
622
+ expected.name = string_series.name
623
+ result = string_series.apply(ops, by_row=by_row)
624
+ tm.assert_series_equal(result, expected)
625
+
626
+
627
+ @pytest.mark.parametrize(
628
+ "ops, expected",
629
+ [
630
+ (
631
+ {"a": lambda x: x},
632
+ Series([1, 2, 3], index=MultiIndex.from_arrays([["a"] * 3, range(3)])),
633
+ ),
634
+ ({"a": lambda x: x.sum()}, Series([6], index=["a"])),
635
+ ],
636
+ )
637
+ def test_apply_dictlike_lambda(ops, by_row, expected):
638
+ # GH53400
639
+ ser = Series([1, 2, 3])
640
+ result = ser.apply(ops, by_row=by_row)
641
+ tm.assert_equal(result, expected)
642
+
643
+
644
+ def test_apply_retains_column_name(by_row):
645
+ # GH 16380
646
+ df = DataFrame({"x": range(3)}, Index(range(3), name="x"))
647
+ result = df.x.apply(lambda x: Series(range(x + 1), Index(range(x + 1), name="y")))
648
+ expected = DataFrame(
649
+ [[0.0, np.nan, np.nan], [0.0, 1.0, np.nan], [0.0, 1.0, 2.0]],
650
+ columns=Index(range(3), name="y"),
651
+ index=Index(range(3), name="x"),
652
+ )
653
+ tm.assert_frame_equal(result, expected)
654
+
655
+
656
+ def test_apply_type():
657
+ # GH 46719
658
+ s = Series([3, "string", float], index=["a", "b", "c"])
659
+ result = s.apply(type)
660
+ expected = Series([int, str, type], index=["a", "b", "c"])
661
+ tm.assert_series_equal(result, expected)
662
+
663
+
664
+ def test_series_apply_unpack_nested_data():
665
+ # GH#55189
666
+ ser = Series([[1, 2, 3], [4, 5, 6, 7]])
667
+ result = ser.apply(lambda x: Series(x))
668
+ expected = DataFrame({0: [1.0, 4.0], 1: [2.0, 5.0], 2: [3.0, 6.0], 3: [np.nan, 7]})
669
+ tm.assert_frame_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_apply_relabeling.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+ import pandas._testing as tm
3
+
4
+
5
+ def test_relabel_no_duplicated_method():
6
+ # this is to test there is no duplicated method used in agg
7
+ df = pd.DataFrame({"A": [1, 2, 1, 2], "B": [1, 2, 3, 4]})
8
+
9
+ result = df["A"].agg(foo="sum")
10
+ expected = df["A"].agg({"foo": "sum"})
11
+ tm.assert_series_equal(result, expected)
12
+
13
+ result = df["B"].agg(foo="min", bar="max")
14
+ expected = df["B"].agg({"foo": "min", "bar": "max"})
15
+ tm.assert_series_equal(result, expected)
16
+
17
+ result = df["B"].agg(foo=sum, bar=min, cat="max")
18
+ expected = df["B"].agg({"foo": sum, "bar": min, "cat": "max"})
19
+ tm.assert_series_equal(result, expected)
20
+
21
+
22
+ def test_relabel_duplicated_method():
23
+ # this is to test with nested renaming, duplicated method can be used
24
+ # if they are assigned with different new names
25
+ df = pd.DataFrame({"A": [1, 2, 1, 2], "B": [1, 2, 3, 4]})
26
+
27
+ result = df["A"].agg(foo="sum", bar="sum")
28
+ expected = pd.Series([6, 6], index=["foo", "bar"], name="A")
29
+ tm.assert_series_equal(result, expected)
30
+
31
+ result = df["B"].agg(foo=min, bar="min")
32
+ expected = pd.Series([1, 1], index=["foo", "bar"], name="B")
33
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_series_transform.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas import (
5
+ DataFrame,
6
+ MultiIndex,
7
+ Series,
8
+ concat,
9
+ )
10
+ import pandas._testing as tm
11
+
12
+
13
+ @pytest.mark.parametrize(
14
+ "args, kwargs, increment",
15
+ [((), {}, 0), ((), {"a": 1}, 1), ((2, 3), {}, 32), ((1,), {"c": 2}, 201)],
16
+ )
17
+ def test_agg_args(args, kwargs, increment):
18
+ # GH 43357
19
+ def f(x, a=0, b=0, c=0):
20
+ return x + a + 10 * b + 100 * c
21
+
22
+ s = Series([1, 2])
23
+ result = s.transform(f, 0, *args, **kwargs)
24
+ expected = s + increment
25
+ tm.assert_series_equal(result, expected)
26
+
27
+
28
+ @pytest.mark.parametrize(
29
+ "ops, names",
30
+ [
31
+ ([np.sqrt], ["sqrt"]),
32
+ ([np.abs, np.sqrt], ["absolute", "sqrt"]),
33
+ (np.array([np.sqrt]), ["sqrt"]),
34
+ (np.array([np.abs, np.sqrt]), ["absolute", "sqrt"]),
35
+ ],
36
+ )
37
+ def test_transform_listlike(string_series, ops, names):
38
+ # GH 35964
39
+ with np.errstate(all="ignore"):
40
+ expected = concat([op(string_series) for op in ops], axis=1)
41
+ expected.columns = names
42
+ result = string_series.transform(ops)
43
+ tm.assert_frame_equal(result, expected)
44
+
45
+
46
+ def test_transform_listlike_func_with_args():
47
+ # GH 50624
48
+
49
+ s = Series([1, 2, 3])
50
+
51
+ def foo1(x, a=1, c=0):
52
+ return x + a + c
53
+
54
+ def foo2(x, b=2, c=0):
55
+ return x + b + c
56
+
57
+ msg = r"foo1\(\) got an unexpected keyword argument 'b'"
58
+ with pytest.raises(TypeError, match=msg):
59
+ s.transform([foo1, foo2], 0, 3, b=3, c=4)
60
+
61
+ result = s.transform([foo1, foo2], 0, 3, c=4)
62
+ expected = DataFrame({"foo1": [8, 9, 10], "foo2": [8, 9, 10]})
63
+ tm.assert_frame_equal(result, expected)
64
+
65
+
66
+ @pytest.mark.parametrize("box", [dict, Series])
67
+ def test_transform_dictlike(string_series, box):
68
+ # GH 35964
69
+ with np.errstate(all="ignore"):
70
+ expected = concat([np.sqrt(string_series), np.abs(string_series)], axis=1)
71
+ expected.columns = ["foo", "bar"]
72
+ result = string_series.transform(box({"foo": np.sqrt, "bar": np.abs}))
73
+ tm.assert_frame_equal(result, expected)
74
+
75
+
76
+ def test_transform_dictlike_mixed():
77
+ # GH 40018 - mix of lists and non-lists in values of a dictionary
78
+ df = Series([1, 4])
79
+ result = df.transform({"b": ["sqrt", "abs"], "c": "sqrt"})
80
+ expected = DataFrame(
81
+ [[1.0, 1, 1.0], [2.0, 4, 2.0]],
82
+ columns=MultiIndex([("b", "c"), ("sqrt", "abs")], [(0, 0, 1), (0, 1, 0)]),
83
+ )
84
+ tm.assert_frame_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/apply/test_str.py ADDED
@@ -0,0 +1,307 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from itertools import chain
2
+ import operator
3
+
4
+ import numpy as np
5
+ import pytest
6
+
7
+ from pandas.compat import (
8
+ WASM,
9
+ )
10
+
11
+ from pandas.core.dtypes.common import is_number
12
+
13
+ from pandas import (
14
+ DataFrame,
15
+ Series,
16
+ )
17
+ import pandas._testing as tm
18
+ from pandas.tests.apply.common import (
19
+ frame_transform_kernels,
20
+ series_transform_kernels,
21
+ )
22
+
23
+
24
+ @pytest.mark.parametrize("func", ["sum", "mean", "min", "max", "std"])
25
+ @pytest.mark.parametrize(
26
+ "kwds",
27
+ [
28
+ pytest.param({}, id="no_kwds"),
29
+ pytest.param({"axis": 1}, id="on_axis"),
30
+ pytest.param({"numeric_only": True}, id="func_kwds"),
31
+ pytest.param({"axis": 1, "numeric_only": True}, id="axis_and_func_kwds"),
32
+ ],
33
+ )
34
+ @pytest.mark.parametrize("how", ["agg", "apply"])
35
+ def test_apply_with_string_funcs(float_frame, func, kwds, how):
36
+ result = getattr(float_frame, how)(func, **kwds)
37
+ expected = getattr(float_frame, func)(**kwds)
38
+ tm.assert_series_equal(result, expected)
39
+
40
+
41
+ def test_with_string_args(datetime_series, all_numeric_reductions):
42
+ result = datetime_series.apply(all_numeric_reductions)
43
+ expected = getattr(datetime_series, all_numeric_reductions)()
44
+ assert result == expected
45
+
46
+
47
+ @pytest.mark.parametrize("op", ["mean", "median", "std", "var"])
48
+ @pytest.mark.parametrize("how", ["agg", "apply"])
49
+ def test_apply_np_reducer(op, how):
50
+ # GH 39116
51
+ float_frame = DataFrame({"a": [1, 2], "b": [3, 4]})
52
+ result = getattr(float_frame, how)(op)
53
+ # pandas ddof defaults to 1, numpy to 0
54
+ kwargs = {"ddof": 1} if op in ("std", "var") else {}
55
+ expected = Series(
56
+ getattr(np, op)(float_frame, axis=0, **kwargs), index=float_frame.columns
57
+ )
58
+ tm.assert_series_equal(result, expected)
59
+
60
+
61
+ @pytest.mark.skipif(WASM, reason="No fp exception support in wasm")
62
+ @pytest.mark.parametrize(
63
+ "op", ["abs", "ceil", "cos", "cumsum", "exp", "log", "sqrt", "square"]
64
+ )
65
+ @pytest.mark.parametrize("how", ["transform", "apply"])
66
+ def test_apply_np_transformer(float_frame, op, how):
67
+ # GH 39116
68
+
69
+ # float_frame will _usually_ have negative values, which will
70
+ # trigger the warning here, but let's put one in just to be sure
71
+ float_frame.iloc[0, 0] = -1.0
72
+ warn = None
73
+ if op in ["log", "sqrt"]:
74
+ warn = RuntimeWarning
75
+
76
+ with tm.assert_produces_warning(warn, check_stacklevel=False):
77
+ # float_frame fixture is defined in conftest.py, so we don't check the
78
+ # stacklevel as otherwise the test would fail.
79
+ result = getattr(float_frame, how)(op)
80
+ expected = getattr(np, op)(float_frame)
81
+ tm.assert_frame_equal(result, expected)
82
+
83
+
84
+ @pytest.mark.parametrize(
85
+ "series, func, expected",
86
+ chain(
87
+ tm.get_cython_table_params(
88
+ Series(dtype=np.float64),
89
+ [
90
+ ("sum", 0),
91
+ ("max", np.nan),
92
+ ("min", np.nan),
93
+ ("all", True),
94
+ ("any", False),
95
+ ("mean", np.nan),
96
+ ("prod", 1),
97
+ ("std", np.nan),
98
+ ("var", np.nan),
99
+ ("median", np.nan),
100
+ ],
101
+ ),
102
+ tm.get_cython_table_params(
103
+ Series([np.nan, 1, 2, 3]),
104
+ [
105
+ ("sum", 6),
106
+ ("max", 3),
107
+ ("min", 1),
108
+ ("all", True),
109
+ ("any", True),
110
+ ("mean", 2),
111
+ ("prod", 6),
112
+ ("std", 1),
113
+ ("var", 1),
114
+ ("median", 2),
115
+ ],
116
+ ),
117
+ tm.get_cython_table_params(
118
+ Series("a b c".split()),
119
+ [
120
+ ("sum", "abc"),
121
+ ("max", "c"),
122
+ ("min", "a"),
123
+ ("all", True),
124
+ ("any", True),
125
+ ],
126
+ ),
127
+ ),
128
+ )
129
+ def test_agg_cython_table_series(series, func, expected):
130
+ # GH21224
131
+ # test reducing functions in
132
+ # pandas.core.base.SelectionMixin._cython_table
133
+ warn = None if isinstance(func, str) else FutureWarning
134
+ with tm.assert_produces_warning(warn, match="is currently using Series.*"):
135
+ result = series.agg(func)
136
+ if is_number(expected):
137
+ assert np.isclose(result, expected, equal_nan=True)
138
+ else:
139
+ assert result == expected
140
+
141
+
142
+ @pytest.mark.parametrize(
143
+ "series, func, expected",
144
+ chain(
145
+ tm.get_cython_table_params(
146
+ Series(dtype=np.float64),
147
+ [
148
+ ("cumprod", Series([], dtype=np.float64)),
149
+ ("cumsum", Series([], dtype=np.float64)),
150
+ ],
151
+ ),
152
+ tm.get_cython_table_params(
153
+ Series([np.nan, 1, 2, 3]),
154
+ [
155
+ ("cumprod", Series([np.nan, 1, 2, 6])),
156
+ ("cumsum", Series([np.nan, 1, 3, 6])),
157
+ ],
158
+ ),
159
+ tm.get_cython_table_params(
160
+ Series("a b c".split()), [("cumsum", Series(["a", "ab", "abc"]))]
161
+ ),
162
+ ),
163
+ )
164
+ def test_agg_cython_table_transform_series(series, func, expected):
165
+ # GH21224
166
+ # test transforming functions in
167
+ # pandas.core.base.SelectionMixin._cython_table (cumprod, cumsum)
168
+ warn = None if isinstance(func, str) else FutureWarning
169
+ with tm.assert_produces_warning(warn, match="is currently using Series.*"):
170
+ result = series.agg(func)
171
+ tm.assert_series_equal(result, expected)
172
+
173
+
174
+ @pytest.mark.parametrize(
175
+ "df, func, expected",
176
+ chain(
177
+ tm.get_cython_table_params(
178
+ DataFrame(),
179
+ [
180
+ ("sum", Series(dtype="float64")),
181
+ ("max", Series(dtype="float64")),
182
+ ("min", Series(dtype="float64")),
183
+ ("all", Series(dtype=bool)),
184
+ ("any", Series(dtype=bool)),
185
+ ("mean", Series(dtype="float64")),
186
+ ("prod", Series(dtype="float64")),
187
+ ("std", Series(dtype="float64")),
188
+ ("var", Series(dtype="float64")),
189
+ ("median", Series(dtype="float64")),
190
+ ],
191
+ ),
192
+ tm.get_cython_table_params(
193
+ DataFrame([[np.nan, 1], [1, 2]]),
194
+ [
195
+ ("sum", Series([1.0, 3])),
196
+ ("max", Series([1.0, 2])),
197
+ ("min", Series([1.0, 1])),
198
+ ("all", Series([True, True])),
199
+ ("any", Series([True, True])),
200
+ ("mean", Series([1, 1.5])),
201
+ ("prod", Series([1.0, 2])),
202
+ ("std", Series([np.nan, 0.707107])),
203
+ ("var", Series([np.nan, 0.5])),
204
+ ("median", Series([1, 1.5])),
205
+ ],
206
+ ),
207
+ ),
208
+ )
209
+ def test_agg_cython_table_frame(df, func, expected, axis):
210
+ # GH 21224
211
+ # test reducing functions in
212
+ # pandas.core.base.SelectionMixin._cython_table
213
+ warn = None if isinstance(func, str) else FutureWarning
214
+ with tm.assert_produces_warning(warn, match="is currently using DataFrame.*"):
215
+ # GH#53425
216
+ result = df.agg(func, axis=axis)
217
+ tm.assert_series_equal(result, expected)
218
+
219
+
220
+ @pytest.mark.parametrize(
221
+ "df, func, expected",
222
+ chain(
223
+ tm.get_cython_table_params(
224
+ DataFrame(), [("cumprod", DataFrame()), ("cumsum", DataFrame())]
225
+ ),
226
+ tm.get_cython_table_params(
227
+ DataFrame([[np.nan, 1], [1, 2]]),
228
+ [
229
+ ("cumprod", DataFrame([[np.nan, 1], [1, 2]])),
230
+ ("cumsum", DataFrame([[np.nan, 1], [1, 3]])),
231
+ ],
232
+ ),
233
+ ),
234
+ )
235
+ def test_agg_cython_table_transform_frame(df, func, expected, axis):
236
+ # GH 21224
237
+ # test transforming functions in
238
+ # pandas.core.base.SelectionMixin._cython_table (cumprod, cumsum)
239
+ if axis in ("columns", 1):
240
+ # operating blockwise doesn't let us preserve dtypes
241
+ expected = expected.astype("float64")
242
+
243
+ warn = None if isinstance(func, str) else FutureWarning
244
+ with tm.assert_produces_warning(warn, match="is currently using DataFrame.*"):
245
+ # GH#53425
246
+ result = df.agg(func, axis=axis)
247
+ tm.assert_frame_equal(result, expected)
248
+
249
+
250
+ @pytest.mark.parametrize("op", series_transform_kernels)
251
+ def test_transform_groupby_kernel_series(request, string_series, op):
252
+ # GH 35964
253
+ if op == "ngroup":
254
+ request.applymarker(
255
+ pytest.mark.xfail(raises=ValueError, reason="ngroup not valid for NDFrame")
256
+ )
257
+ args = [0.0] if op == "fillna" else []
258
+ ones = np.ones(string_series.shape[0])
259
+
260
+ warn = FutureWarning if op == "fillna" else None
261
+ msg = "SeriesGroupBy.fillna is deprecated"
262
+ with tm.assert_produces_warning(warn, match=msg):
263
+ expected = string_series.groupby(ones).transform(op, *args)
264
+ result = string_series.transform(op, 0, *args)
265
+ tm.assert_series_equal(result, expected)
266
+
267
+
268
+ @pytest.mark.parametrize("op", frame_transform_kernels)
269
+ def test_transform_groupby_kernel_frame(request, float_frame, op):
270
+ if op == "ngroup":
271
+ request.applymarker(
272
+ pytest.mark.xfail(raises=ValueError, reason="ngroup not valid for NDFrame")
273
+ )
274
+
275
+ # GH 35964
276
+
277
+ args = [0.0] if op == "fillna" else []
278
+ ones = np.ones(float_frame.shape[0])
279
+ gb = float_frame.groupby(ones)
280
+
281
+ warn = FutureWarning if op == "fillna" else None
282
+ op_msg = "DataFrameGroupBy.fillna is deprecated"
283
+ with tm.assert_produces_warning(warn, match=op_msg):
284
+ expected = gb.transform(op, *args)
285
+
286
+ result = float_frame.transform(op, 0, *args)
287
+ tm.assert_frame_equal(result, expected)
288
+
289
+ # same thing, but ensuring we have multiple blocks
290
+ assert "E" not in float_frame.columns
291
+ float_frame["E"] = float_frame["A"].copy()
292
+ assert len(float_frame._mgr.blocks) > 1
293
+
294
+ ones = np.ones(float_frame.shape[0])
295
+ gb2 = float_frame.groupby(ones)
296
+ expected2 = gb2.transform(op, *args)
297
+ result2 = float_frame.transform(op, 0, *args)
298
+ tm.assert_frame_equal(result2, expected2)
299
+
300
+
301
+ @pytest.mark.parametrize("method", ["abs", "shift", "pct_change", "cumsum", "rank"])
302
+ def test_transform_method_name(method):
303
+ # GH 19760
304
+ df = DataFrame({"A": [-1, 2]})
305
+ result = df.transform(method)
306
+ expected = operator.methodcaller(method)(df)
307
+ tm.assert_frame_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/__init__.py ADDED
File without changes
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/common.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Assertion helpers for arithmetic tests.
3
+ """
4
+
5
+ import numpy as np
6
+ import pytest
7
+
8
+ from pandas import (
9
+ DataFrame,
10
+ Index,
11
+ Series,
12
+ array,
13
+ )
14
+ import pandas._testing as tm
15
+ from pandas.core.arrays import (
16
+ BooleanArray,
17
+ NumpyExtensionArray,
18
+ )
19
+
20
+
21
+ def assert_cannot_add(left, right, msg="cannot add"):
22
+ """
23
+ Helper function to assert that two objects cannot be added.
24
+
25
+ Parameters
26
+ ----------
27
+ left : object
28
+ The first operand.
29
+ right : object
30
+ The second operand.
31
+ msg : str, default "cannot add"
32
+ The error message expected in the TypeError.
33
+ """
34
+ with pytest.raises(TypeError, match=msg):
35
+ left + right
36
+ with pytest.raises(TypeError, match=msg):
37
+ right + left
38
+
39
+
40
+ def assert_invalid_addsub_type(left, right, msg=None):
41
+ """
42
+ Helper function to assert that two objects can
43
+ neither be added nor subtracted.
44
+
45
+ Parameters
46
+ ----------
47
+ left : object
48
+ The first operand.
49
+ right : object
50
+ The second operand.
51
+ msg : str or None, default None
52
+ The error message expected in the TypeError.
53
+ """
54
+ with pytest.raises(TypeError, match=msg):
55
+ left + right
56
+ with pytest.raises(TypeError, match=msg):
57
+ right + left
58
+ with pytest.raises(TypeError, match=msg):
59
+ left - right
60
+ with pytest.raises(TypeError, match=msg):
61
+ right - left
62
+
63
+
64
+ def get_upcast_box(left, right, is_cmp: bool = False):
65
+ """
66
+ Get the box to use for 'expected' in an arithmetic or comparison operation.
67
+
68
+ Parameters
69
+ left : Any
70
+ right : Any
71
+ is_cmp : bool, default False
72
+ Whether the operation is a comparison method.
73
+ """
74
+
75
+ if isinstance(left, DataFrame) or isinstance(right, DataFrame):
76
+ return DataFrame
77
+ if isinstance(left, Series) or isinstance(right, Series):
78
+ if is_cmp and isinstance(left, Index):
79
+ # Index does not defer for comparisons
80
+ return np.array
81
+ return Series
82
+ if isinstance(left, Index) or isinstance(right, Index):
83
+ if is_cmp:
84
+ return np.array
85
+ return Index
86
+ return tm.to_array
87
+
88
+
89
+ def assert_invalid_comparison(left, right, box):
90
+ """
91
+ Assert that comparison operations with mismatched types behave correctly.
92
+
93
+ Parameters
94
+ ----------
95
+ left : np.ndarray, ExtensionArray, Index, or Series
96
+ right : object
97
+ box : {pd.DataFrame, pd.Series, pd.Index, pd.array, tm.to_array}
98
+ """
99
+ # Not for tznaive-tzaware comparison
100
+
101
+ # Note: not quite the same as how we do this for tm.box_expected
102
+ xbox = box if box not in [Index, array] else np.array
103
+
104
+ def xbox2(x):
105
+ # Eventually we'd like this to be tighter, but for now we'll
106
+ # just exclude NumpyExtensionArray[bool]
107
+ if isinstance(x, NumpyExtensionArray):
108
+ return x._ndarray
109
+ if isinstance(x, BooleanArray):
110
+ # NB: we are assuming no pd.NAs for now
111
+ return x.astype(bool)
112
+ return x
113
+
114
+ result = xbox2(left == right)
115
+ expected = xbox(np.zeros(result.shape, dtype=np.bool_))
116
+
117
+ tm.assert_equal(result, expected)
118
+
119
+ result = xbox2(right == left)
120
+ tm.assert_equal(result, xbox(expected))
121
+
122
+ result = xbox2(left != right)
123
+ tm.assert_equal(result, ~expected)
124
+
125
+ result = xbox2(right != left)
126
+ tm.assert_equal(result, xbox(~expected))
127
+
128
+ msg = "|".join(
129
+ [
130
+ "Invalid comparison between",
131
+ "Cannot compare type",
132
+ "not supported between",
133
+ "invalid type promotion",
134
+ (
135
+ # GH#36706 npdev 1.20.0 2020-09-28
136
+ r"The DTypes <class 'numpy.dtype\[datetime64\]'> and "
137
+ r"<class 'numpy.dtype\[int64\]'> do not have a common DType. "
138
+ "For example they cannot be stored in a single array unless the "
139
+ "dtype is `object`."
140
+ ),
141
+ ]
142
+ )
143
+ with pytest.raises(TypeError, match=msg):
144
+ left < right
145
+ with pytest.raises(TypeError, match=msg):
146
+ left <= right
147
+ with pytest.raises(TypeError, match=msg):
148
+ left > right
149
+ with pytest.raises(TypeError, match=msg):
150
+ left >= right
151
+ with pytest.raises(TypeError, match=msg):
152
+ right < left
153
+ with pytest.raises(TypeError, match=msg):
154
+ right <= left
155
+ with pytest.raises(TypeError, match=msg):
156
+ right > left
157
+ with pytest.raises(TypeError, match=msg):
158
+ right >= left
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/conftest.py ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ from pandas import Index
6
+
7
+
8
+ @pytest.fixture(params=[1, np.array(1, dtype=np.int64)])
9
+ def one(request):
10
+ """
11
+ Several variants of integer value 1. The zero-dim integer array
12
+ behaves like an integer.
13
+
14
+ This fixture can be used to check that datetimelike indexes handle
15
+ addition and subtraction of integers and zero-dimensional arrays
16
+ of integers.
17
+
18
+ Examples
19
+ --------
20
+ dti = pd.date_range('2016-01-01', periods=2, freq='h')
21
+ dti
22
+ DatetimeIndex(['2016-01-01 00:00:00', '2016-01-01 01:00:00'],
23
+ dtype='datetime64[ns]', freq='h')
24
+ dti + one
25
+ DatetimeIndex(['2016-01-01 01:00:00', '2016-01-01 02:00:00'],
26
+ dtype='datetime64[ns]', freq='h')
27
+ """
28
+ return request.param
29
+
30
+
31
+ zeros = [
32
+ box_cls([0] * 5, dtype=dtype)
33
+ for box_cls in [Index, np.array, pd.array]
34
+ for dtype in [np.int64, np.uint64, np.float64]
35
+ ]
36
+ zeros.extend([box_cls([-0.0] * 5, dtype=np.float64) for box_cls in [Index, np.array]])
37
+ zeros.extend([np.array(0, dtype=dtype) for dtype in [np.int64, np.uint64, np.float64]])
38
+ zeros.extend([np.array(-0.0, dtype=np.float64)])
39
+ zeros.extend([0, 0.0, -0.0])
40
+
41
+
42
+ @pytest.fixture(params=zeros)
43
+ def zero(request):
44
+ """
45
+ Several types of scalar zeros and length 5 vectors of zeros.
46
+
47
+ This fixture can be used to check that numeric-dtype indexes handle
48
+ division by any zero numeric-dtype.
49
+
50
+ Uses vector of length 5 for broadcasting with `numeric_idx` fixture,
51
+ which creates numeric-dtype vectors also of length 5.
52
+
53
+ Examples
54
+ --------
55
+ arr = RangeIndex(5)
56
+ arr / zeros
57
+ Index([nan, inf, inf, inf, inf], dtype='float64')
58
+ """
59
+ return request.param
60
+
61
+
62
+ # ------------------------------------------------------------------
63
+ # Scalar Fixtures
64
+
65
+
66
+ @pytest.fixture(
67
+ params=[
68
+ pd.Timedelta("10m7s").to_pytimedelta(),
69
+ pd.Timedelta("10m7s"),
70
+ pd.Timedelta("10m7s").to_timedelta64(),
71
+ ],
72
+ ids=lambda x: type(x).__name__,
73
+ )
74
+ def scalar_td(request):
75
+ """
76
+ Several variants of Timedelta scalars representing 10 minutes and 7 seconds.
77
+ """
78
+ return request.param
79
+
80
+
81
+ @pytest.fixture(
82
+ params=[
83
+ pd.offsets.Day(3),
84
+ pd.offsets.Hour(72),
85
+ pd.Timedelta(days=3).to_pytimedelta(),
86
+ pd.Timedelta("72:00:00"),
87
+ np.timedelta64(3, "D"),
88
+ np.timedelta64(72, "h"),
89
+ ],
90
+ ids=lambda x: type(x).__name__,
91
+ )
92
+ def three_days(request):
93
+ """
94
+ Several timedelta-like and DateOffset objects that each represent
95
+ a 3-day timedelta
96
+ """
97
+ return request.param
98
+
99
+
100
+ @pytest.fixture(
101
+ params=[
102
+ pd.offsets.Hour(2),
103
+ pd.offsets.Minute(120),
104
+ pd.Timedelta(hours=2).to_pytimedelta(),
105
+ pd.Timedelta(seconds=2 * 3600),
106
+ np.timedelta64(2, "h"),
107
+ np.timedelta64(120, "m"),
108
+ ],
109
+ ids=lambda x: type(x).__name__,
110
+ )
111
+ def two_hours(request):
112
+ """
113
+ Several timedelta-like and DateOffset objects that each represent
114
+ a 2-hour timedelta
115
+ """
116
+ return request.param
117
+
118
+
119
+ _common_mismatch = [
120
+ pd.offsets.YearBegin(2),
121
+ pd.offsets.MonthBegin(1),
122
+ pd.offsets.Minute(),
123
+ ]
124
+
125
+
126
+ @pytest.fixture(
127
+ params=[
128
+ np.timedelta64(4, "h"),
129
+ pd.Timedelta(hours=23).to_pytimedelta(),
130
+ pd.Timedelta("23:00:00"),
131
+ *_common_mismatch,
132
+ ]
133
+ )
134
+ def not_daily(request):
135
+ """
136
+ Several timedelta-like and DateOffset instances that are _not_
137
+ compatible with Daily frequencies.
138
+ """
139
+ return request.param
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_array_ops.py ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import operator
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ from pandas.core.dtypes.missing import isna
7
+
8
+ import pandas._testing as tm
9
+ from pandas.core.ops.array_ops import (
10
+ comparison_op,
11
+ na_logical_op,
12
+ )
13
+
14
+
15
+ def test_na_logical_op_2d():
16
+ left = np.arange(8).reshape(4, 2)
17
+ right = left.astype(object)
18
+ right[0, 0] = np.nan
19
+
20
+ # Check that we fall back to the vec_binop branch
21
+ with pytest.raises(TypeError, match="unsupported operand type"):
22
+ operator.or_(left, right)
23
+
24
+ result = na_logical_op(left, right, operator.or_)
25
+ expected = right
26
+ tm.assert_numpy_array_equal(result, expected)
27
+
28
+
29
+ def test_object_comparison_2d():
30
+ left = np.arange(9).reshape(3, 3).astype(object)
31
+ right = left.T
32
+
33
+ result = comparison_op(left, right, operator.eq)
34
+ expected = np.eye(3).astype(bool)
35
+ tm.assert_numpy_array_equal(result, expected)
36
+
37
+ # Ensure that cython doesn't raise on non-writeable arg, which
38
+ # we can get from np.broadcast_to
39
+ right.flags.writeable = False
40
+ result = comparison_op(left, right, operator.ne)
41
+ tm.assert_numpy_array_equal(result, ~expected)
42
+
43
+
44
+ @pytest.mark.parametrize("rvalues", [1, [1, 1, 1], np.nan, None])
45
+ @pytest.mark.parametrize(
46
+ "op", [operator.eq, operator.ne, operator.lt, operator.le, operator.gt, operator.ge]
47
+ )
48
+ def test_comparison_for_subclasses(rvalues, op):
49
+ # GH#63205 Ensure subclasses of ndarray are correctly handled in comparison_op
50
+ # Define a custom ndarray subclass
51
+ class TestArray(np.ndarray):
52
+ def __new__(cls, input_array):
53
+ return np.asarray(input_array).view(cls)
54
+
55
+ def __array_finalize__(self, obj) -> None:
56
+ self._is_test_array = True
57
+
58
+ def expected_with_na_handling(lvalues, rvalues, op):
59
+ # Similar to comparison_op, handle zerodim arrays with na value separately
60
+ if (rvalues.ndim == 0) and isna(rvalues.item()):
61
+ # numpy does not like comparisons vs None
62
+ if op is operator.ne:
63
+ return np.ones(lvalues.shape, dtype=bool)
64
+ else:
65
+ return np.zeros(lvalues.shape, dtype=bool)
66
+ return op(lvalues, rvalues)
67
+
68
+ # Define test data
69
+ lvalues = [1, 2, 3]
70
+
71
+ # Test with both ndarray and TestArray
72
+ result = comparison_op(np.array(lvalues), np.array(rvalues), op)
73
+ expected = expected_with_na_handling(np.array(lvalues), np.array(rvalues), op)
74
+ tm.assert_numpy_array_equal(result, expected)
75
+
76
+ result = comparison_op(TestArray(lvalues), TestArray(rvalues), op)
77
+ expected = expected_with_na_handling(TestArray(lvalues), TestArray(rvalues), op)
78
+ tm.assert_numpy_array_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_bool.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pytest
2
+
3
+ from pandas import (
4
+ DataFrame,
5
+ Series,
6
+ )
7
+ import pandas._testing as tm
8
+
9
+
10
+ def test_divmod_bool_raises(box_with_array):
11
+ # GH#46043 // raises, so divmod should too
12
+ ser = Series([True, False])
13
+ obj = tm.box_expected(ser, box_with_array)
14
+
15
+ msg = "operator 'floordiv' not implemented for bool dtypes"
16
+ with pytest.raises(NotImplementedError, match=msg):
17
+ obj // obj
18
+
19
+ if box_with_array is DataFrame:
20
+ msg = "operator 'floordiv' not implemented for bool dtypes"
21
+ else:
22
+ msg = "operator 'divmod' not implemented for bool dtypes"
23
+ with pytest.raises(NotImplementedError, match=msg):
24
+ divmod(obj, obj)
25
+
26
+ # go through __rdivmod__
27
+ with pytest.raises(NotImplementedError, match=msg):
28
+ divmod(True, obj)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_categorical.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ from pandas import (
4
+ Categorical,
5
+ Series,
6
+ )
7
+ import pandas._testing as tm
8
+
9
+
10
+ class TestCategoricalComparisons:
11
+ def test_categorical_nan_equality(self):
12
+ cat = Series(Categorical(["a", "b", "c", np.nan]))
13
+ expected = Series([True, True, True, False])
14
+ result = cat == cat
15
+ tm.assert_series_equal(result, expected)
16
+
17
+ def test_categorical_tuple_equality(self):
18
+ # GH 18050
19
+ ser = Series([(0, 0), (0, 1), (0, 0), (1, 0), (1, 1)])
20
+ expected = Series([True, False, True, False, False])
21
+ result = ser == (0, 0)
22
+ tm.assert_series_equal(result, expected)
23
+
24
+ result = ser.astype("category") == (0, 0)
25
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_datetime64.py ADDED
@@ -0,0 +1,2499 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arithmetic tests for DataFrame/Series/Index/Array classes that should
2
+ # behave identically.
3
+ # Specifically for datetime64 and datetime64tz dtypes
4
+ from datetime import (
5
+ datetime,
6
+ time,
7
+ timedelta,
8
+ timezone,
9
+ )
10
+ from itertools import (
11
+ product,
12
+ )
13
+ import operator
14
+
15
+ import numpy as np
16
+ import pytest
17
+
18
+ from pandas._libs.tslibs.conversion import localize_pydatetime
19
+ from pandas._libs.tslibs.offsets import shift_months
20
+
21
+ import pandas as pd
22
+ from pandas import (
23
+ DateOffset,
24
+ DatetimeIndex,
25
+ NaT,
26
+ Period,
27
+ Series,
28
+ Timedelta,
29
+ TimedeltaIndex,
30
+ Timestamp,
31
+ date_range,
32
+ )
33
+ import pandas._testing as tm
34
+ from pandas.core import roperator
35
+ from pandas.tests.arithmetic.common import (
36
+ assert_cannot_add,
37
+ assert_invalid_addsub_type,
38
+ assert_invalid_comparison,
39
+ get_upcast_box,
40
+ )
41
+
42
+ # ------------------------------------------------------------------
43
+ # Comparisons
44
+
45
+
46
+ class TestDatetime64ArrayLikeComparisons:
47
+ # Comparison tests for datetime64 vectors fully parametrized over
48
+ # DataFrame/Series/DatetimeIndex/DatetimeArray. Ideally all comparison
49
+ # tests will eventually end up here.
50
+
51
+ def test_compare_zerodim(self, tz_naive_fixture, box_with_array):
52
+ # Test comparison with zero-dimensional array is unboxed
53
+ tz = tz_naive_fixture
54
+ box = box_with_array
55
+ dti = date_range("20130101", periods=3, tz=tz)
56
+
57
+ other = np.array(dti.to_numpy()[0])
58
+
59
+ dtarr = tm.box_expected(dti, box)
60
+ xbox = get_upcast_box(dtarr, other, True)
61
+ result = dtarr <= other
62
+ expected = np.array([True, False, False])
63
+ expected = tm.box_expected(expected, xbox)
64
+ tm.assert_equal(result, expected)
65
+
66
+ @pytest.mark.parametrize(
67
+ "other",
68
+ [
69
+ "foo",
70
+ -1,
71
+ 99,
72
+ 4.0,
73
+ object(),
74
+ timedelta(days=2),
75
+ # GH#19800, GH#19301 datetime.date comparison raises to
76
+ # match DatetimeIndex/Timestamp. This also matches the behavior
77
+ # of stdlib datetime.datetime
78
+ datetime(2001, 1, 1).date(),
79
+ # GH#19301 None and NaN are *not* cast to NaT for comparisons
80
+ None,
81
+ np.nan,
82
+ ],
83
+ )
84
+ def test_dt64arr_cmp_scalar_invalid(self, other, tz_naive_fixture, box_with_array):
85
+ # GH#22074, GH#15966
86
+ tz = tz_naive_fixture
87
+
88
+ rng = date_range("1/1/2000", periods=10, tz=tz)
89
+ dtarr = tm.box_expected(rng, box_with_array)
90
+ assert_invalid_comparison(dtarr, other, box_with_array)
91
+
92
+ @pytest.mark.parametrize(
93
+ "other",
94
+ [
95
+ # GH#4968 invalid date/int comparisons
96
+ list(range(10)),
97
+ np.arange(10),
98
+ np.arange(10).astype(np.float32),
99
+ np.arange(10).astype(object),
100
+ pd.timedelta_range("1ns", periods=10).array,
101
+ np.array(pd.timedelta_range("1ns", periods=10)),
102
+ list(pd.timedelta_range("1ns", periods=10)),
103
+ pd.timedelta_range("1 Day", periods=10).astype(object),
104
+ pd.period_range("1971-01-01", freq="D", periods=10).array,
105
+ pd.period_range("1971-01-01", freq="D", periods=10).astype(object),
106
+ ],
107
+ )
108
+ def test_dt64arr_cmp_arraylike_invalid(
109
+ self, other, tz_naive_fixture, box_with_array
110
+ ):
111
+ tz = tz_naive_fixture
112
+
113
+ dta = date_range("2000-01-01", freq="ns", periods=10, tz=tz)._data
114
+ obj = tm.box_expected(dta, box_with_array)
115
+ assert_invalid_comparison(obj, other, box_with_array)
116
+
117
+ def test_dt64arr_cmp_mixed_invalid(self, tz_naive_fixture):
118
+ tz = tz_naive_fixture
119
+
120
+ dta = date_range("2000-01-01", freq="h", periods=5, tz=tz)._data
121
+ other = np.array([0, 1, 2, dta[3], Timedelta(days=1)])
122
+ result = dta == other
123
+ expected = np.array([False, False, False, True, False])
124
+ tm.assert_numpy_array_equal(result, expected)
125
+
126
+ result = dta != other
127
+ tm.assert_numpy_array_equal(result, ~expected)
128
+
129
+ msg = "Invalid comparison between|Cannot compare type|not supported between"
130
+ with pytest.raises(TypeError, match=msg):
131
+ dta < other
132
+ with pytest.raises(TypeError, match=msg):
133
+ dta > other
134
+ with pytest.raises(TypeError, match=msg):
135
+ dta <= other
136
+ with pytest.raises(TypeError, match=msg):
137
+ dta >= other
138
+
139
+ def test_dt64arr_nat_comparison(self, tz_naive_fixture, box_with_array):
140
+ # GH#22242, GH#22163 DataFrame considered NaT == ts incorrectly
141
+ tz = tz_naive_fixture
142
+ box = box_with_array
143
+
144
+ ts = Timestamp("2021-01-01", tz=tz)
145
+ ser = Series([ts, NaT])
146
+
147
+ obj = tm.box_expected(ser, box)
148
+ xbox = get_upcast_box(obj, ts, True)
149
+
150
+ expected = Series([True, False], dtype=np.bool_)
151
+ expected = tm.box_expected(expected, xbox)
152
+
153
+ result = obj == ts
154
+ tm.assert_equal(result, expected)
155
+
156
+
157
+ class TestDatetime64SeriesComparison:
158
+ # TODO: moved from tests.series.test_operators; needs cleanup
159
+
160
+ @pytest.mark.parametrize(
161
+ "pair",
162
+ [
163
+ (
164
+ [Timestamp("2011-01-01"), NaT, Timestamp("2011-01-03")],
165
+ [NaT, NaT, Timestamp("2011-01-03")],
166
+ ),
167
+ (
168
+ [Timedelta("1 days"), NaT, Timedelta("3 days")],
169
+ [NaT, NaT, Timedelta("3 days")],
170
+ ),
171
+ (
172
+ [Period("2011-01", freq="M"), NaT, Period("2011-03", freq="M")],
173
+ [NaT, NaT, Period("2011-03", freq="M")],
174
+ ),
175
+ ],
176
+ )
177
+ @pytest.mark.parametrize("reverse", [True, False])
178
+ @pytest.mark.parametrize("dtype", [None, object])
179
+ @pytest.mark.parametrize(
180
+ "op, expected",
181
+ [
182
+ (operator.eq, [False, False, True]),
183
+ (operator.ne, [True, True, False]),
184
+ (operator.lt, [False, False, False]),
185
+ (operator.gt, [False, False, False]),
186
+ (operator.ge, [False, False, True]),
187
+ (operator.le, [False, False, True]),
188
+ ],
189
+ )
190
+ def test_nat_comparisons(
191
+ self,
192
+ dtype,
193
+ index_or_series,
194
+ reverse,
195
+ pair,
196
+ op,
197
+ expected,
198
+ ):
199
+ box = index_or_series
200
+ lhs, rhs = pair
201
+ if reverse:
202
+ # add lhs / rhs switched data
203
+ lhs, rhs = rhs, lhs
204
+
205
+ left = Series(lhs, dtype=dtype)
206
+ right = box(rhs, dtype=dtype)
207
+
208
+ result = op(left, right)
209
+
210
+ tm.assert_series_equal(result, Series(expected))
211
+
212
+ @pytest.mark.parametrize(
213
+ "data",
214
+ [
215
+ [Timestamp("2011-01-01"), NaT, Timestamp("2011-01-03")],
216
+ [Timedelta("1 days"), NaT, Timedelta("3 days")],
217
+ [Period("2011-01", freq="M"), NaT, Period("2011-03", freq="M")],
218
+ ],
219
+ )
220
+ @pytest.mark.parametrize("dtype", [None, object])
221
+ def test_nat_comparisons_scalar(self, dtype, data, box_with_array):
222
+ box = box_with_array
223
+
224
+ left = Series(data, dtype=dtype)
225
+ left = tm.box_expected(left, box)
226
+ xbox = get_upcast_box(left, NaT, True)
227
+
228
+ expected = [False, False, False]
229
+ expected = tm.box_expected(expected, xbox)
230
+ if box is pd.array and dtype is object:
231
+ expected = pd.array(expected, dtype="bool")
232
+
233
+ tm.assert_equal(left == NaT, expected)
234
+ tm.assert_equal(NaT == left, expected)
235
+
236
+ expected = [True, True, True]
237
+ expected = tm.box_expected(expected, xbox)
238
+ if box is pd.array and dtype is object:
239
+ expected = pd.array(expected, dtype="bool")
240
+ tm.assert_equal(left != NaT, expected)
241
+ tm.assert_equal(NaT != left, expected)
242
+
243
+ expected = [False, False, False]
244
+ expected = tm.box_expected(expected, xbox)
245
+ if box is pd.array and dtype is object:
246
+ expected = pd.array(expected, dtype="bool")
247
+ tm.assert_equal(left < NaT, expected)
248
+ tm.assert_equal(NaT > left, expected)
249
+ tm.assert_equal(left <= NaT, expected)
250
+ tm.assert_equal(NaT >= left, expected)
251
+
252
+ tm.assert_equal(left > NaT, expected)
253
+ tm.assert_equal(NaT < left, expected)
254
+ tm.assert_equal(left >= NaT, expected)
255
+ tm.assert_equal(NaT <= left, expected)
256
+
257
+ @pytest.mark.parametrize("val", [datetime(2000, 1, 4), datetime(2000, 1, 5)])
258
+ def test_series_comparison_scalars(self, val):
259
+ series = Series(date_range("1/1/2000", periods=10))
260
+
261
+ result = series > val
262
+ expected = Series([x > val for x in series])
263
+ tm.assert_series_equal(result, expected)
264
+
265
+ @pytest.mark.parametrize(
266
+ "left,right", [("lt", "gt"), ("le", "ge"), ("eq", "eq"), ("ne", "ne")]
267
+ )
268
+ def test_timestamp_compare_series(self, left, right):
269
+ # see gh-4982
270
+ # Make sure we can compare Timestamps on the right AND left hand side.
271
+ ser = Series(date_range("20010101", periods=10), name="dates")
272
+ s_nat = ser.copy(deep=True)
273
+
274
+ ser[0] = Timestamp("nat")
275
+ ser[3] = Timestamp("nat")
276
+
277
+ left_f = getattr(operator, left)
278
+ right_f = getattr(operator, right)
279
+
280
+ # No NaT
281
+ expected = left_f(ser, Timestamp("20010109"))
282
+ result = right_f(Timestamp("20010109"), ser)
283
+ tm.assert_series_equal(result, expected)
284
+
285
+ # NaT
286
+ expected = left_f(ser, Timestamp("nat"))
287
+ result = right_f(Timestamp("nat"), ser)
288
+ tm.assert_series_equal(result, expected)
289
+
290
+ # Compare to Timestamp with series containing NaT
291
+ expected = left_f(s_nat, Timestamp("20010109"))
292
+ result = right_f(Timestamp("20010109"), s_nat)
293
+ tm.assert_series_equal(result, expected)
294
+
295
+ # Compare to NaT with series containing NaT
296
+ expected = left_f(s_nat, NaT)
297
+ result = right_f(NaT, s_nat)
298
+ tm.assert_series_equal(result, expected)
299
+
300
+ def test_dt64arr_timestamp_equality(self, box_with_array):
301
+ # GH#11034
302
+ box = box_with_array
303
+
304
+ ser = Series([Timestamp("2000-01-29 01:59:00"), Timestamp("2000-01-30"), NaT])
305
+ ser = tm.box_expected(ser, box)
306
+ xbox = get_upcast_box(ser, ser, True)
307
+
308
+ result = ser != ser
309
+ expected = tm.box_expected([False, False, True], xbox)
310
+ tm.assert_equal(result, expected)
311
+
312
+ if box is pd.DataFrame:
313
+ # alignment for frame vs series comparisons deprecated
314
+ # in GH#46795 enforced 2.0
315
+ with pytest.raises(ValueError, match="not aligned"):
316
+ ser != ser[0]
317
+
318
+ else:
319
+ result = ser != ser[0]
320
+ expected = tm.box_expected([False, True, True], xbox)
321
+ tm.assert_equal(result, expected)
322
+
323
+ if box is pd.DataFrame:
324
+ # alignment for frame vs series comparisons deprecated
325
+ # in GH#46795 enforced 2.0
326
+ with pytest.raises(ValueError, match="not aligned"):
327
+ ser != ser[2]
328
+ else:
329
+ result = ser != ser[2]
330
+ expected = tm.box_expected([True, True, True], xbox)
331
+ tm.assert_equal(result, expected)
332
+
333
+ result = ser == ser
334
+ expected = tm.box_expected([True, True, False], xbox)
335
+ tm.assert_equal(result, expected)
336
+
337
+ if box is pd.DataFrame:
338
+ # alignment for frame vs series comparisons deprecated
339
+ # in GH#46795 enforced 2.0
340
+ with pytest.raises(ValueError, match="not aligned"):
341
+ ser == ser[0]
342
+ else:
343
+ result = ser == ser[0]
344
+ expected = tm.box_expected([True, False, False], xbox)
345
+ tm.assert_equal(result, expected)
346
+
347
+ if box is pd.DataFrame:
348
+ # alignment for frame vs series comparisons deprecated
349
+ # in GH#46795 enforced 2.0
350
+ with pytest.raises(ValueError, match="not aligned"):
351
+ ser == ser[2]
352
+ else:
353
+ result = ser == ser[2]
354
+ expected = tm.box_expected([False, False, False], xbox)
355
+ tm.assert_equal(result, expected)
356
+
357
+ @pytest.mark.parametrize(
358
+ "datetimelike",
359
+ [
360
+ Timestamp("20130101"),
361
+ datetime(2013, 1, 1),
362
+ np.datetime64("2013-01-01T00:00", "ns"),
363
+ ],
364
+ )
365
+ @pytest.mark.parametrize(
366
+ "op,expected",
367
+ [
368
+ (operator.lt, [True, False, False, False]),
369
+ (operator.le, [True, True, False, False]),
370
+ (operator.eq, [False, True, False, False]),
371
+ (operator.gt, [False, False, False, True]),
372
+ ],
373
+ )
374
+ def test_dt64_compare_datetime_scalar(self, datetimelike, op, expected):
375
+ # GH#17965, test for ability to compare datetime64[ns] columns
376
+ # to datetimelike
377
+ ser = Series(
378
+ [
379
+ Timestamp("20120101"),
380
+ Timestamp("20130101"),
381
+ np.nan,
382
+ Timestamp("20130103"),
383
+ ],
384
+ name="A",
385
+ )
386
+ result = op(ser, datetimelike)
387
+ expected = Series(expected, name="A")
388
+ tm.assert_series_equal(result, expected)
389
+
390
+ def test_ts_series_numpy_maximum(self):
391
+ # GH#50864, test numpy.maximum does not fail
392
+ # given a TimeStamp and Series(with dtype datetime64) comparison
393
+ ts = Timestamp("2024-07-01")
394
+ ts_series = Series(
395
+ ["2024-06-01", "2024-07-01", "2024-08-01"],
396
+ dtype="datetime64[us]",
397
+ )
398
+
399
+ expected = Series(
400
+ ["2024-07-01", "2024-07-01", "2024-08-01"],
401
+ dtype="datetime64[us]",
402
+ )
403
+
404
+ tm.assert_series_equal(expected, np.maximum(ts, ts_series))
405
+
406
+
407
+ class TestDatetimeIndexComparisons:
408
+ # TODO: moved from tests.indexes.test_base; parametrize and de-duplicate
409
+ def test_comparators(self, comparison_op):
410
+ index = date_range("2020-01-01", periods=10)
411
+ element = index[len(index) // 2]
412
+ element = Timestamp(element).to_datetime64()
413
+
414
+ arr = np.array(index)
415
+ arr_result = comparison_op(arr, element)
416
+ index_result = comparison_op(index, element)
417
+
418
+ assert isinstance(index_result, np.ndarray)
419
+ tm.assert_numpy_array_equal(arr_result, index_result)
420
+
421
+ @pytest.mark.parametrize(
422
+ "other",
423
+ [datetime(2016, 1, 1), Timestamp("2016-01-01"), np.datetime64("2016-01-01")],
424
+ )
425
+ def test_dti_cmp_datetimelike(self, other, tz_naive_fixture):
426
+ tz = tz_naive_fixture
427
+ dti = date_range("2016-01-01", periods=2, tz=tz)
428
+ if tz is not None:
429
+ if isinstance(other, np.datetime64):
430
+ pytest.skip(f"{type(other).__name__} is not tz aware")
431
+ other = localize_pydatetime(other, dti.tzinfo)
432
+
433
+ result = dti == other
434
+ expected = np.array([True, False])
435
+ tm.assert_numpy_array_equal(result, expected)
436
+
437
+ result = dti > other
438
+ expected = np.array([False, True])
439
+ tm.assert_numpy_array_equal(result, expected)
440
+
441
+ result = dti >= other
442
+ expected = np.array([True, True])
443
+ tm.assert_numpy_array_equal(result, expected)
444
+
445
+ result = dti < other
446
+ expected = np.array([False, False])
447
+ tm.assert_numpy_array_equal(result, expected)
448
+
449
+ result = dti <= other
450
+ expected = np.array([True, False])
451
+ tm.assert_numpy_array_equal(result, expected)
452
+
453
+ @pytest.mark.parametrize("dtype", [None, object])
454
+ def test_dti_cmp_nat(self, dtype, box_with_array):
455
+ left = DatetimeIndex([Timestamp("2011-01-01"), NaT, Timestamp("2011-01-03")])
456
+ right = DatetimeIndex([NaT, NaT, Timestamp("2011-01-03")])
457
+
458
+ left = tm.box_expected(left, box_with_array)
459
+ right = tm.box_expected(right, box_with_array)
460
+ xbox = get_upcast_box(left, right, True)
461
+
462
+ lhs, rhs = left, right
463
+ if dtype is object:
464
+ lhs, rhs = left.astype(object), right.astype(object)
465
+
466
+ result = rhs == lhs
467
+ expected = np.array([False, False, True])
468
+ expected = tm.box_expected(expected, xbox)
469
+ tm.assert_equal(result, expected)
470
+
471
+ result = lhs != rhs
472
+ expected = np.array([True, True, False])
473
+ expected = tm.box_expected(expected, xbox)
474
+ tm.assert_equal(result, expected)
475
+
476
+ expected = np.array([False, False, False])
477
+ expected = tm.box_expected(expected, xbox)
478
+ tm.assert_equal(lhs == NaT, expected)
479
+ tm.assert_equal(NaT == rhs, expected)
480
+
481
+ expected = np.array([True, True, True])
482
+ expected = tm.box_expected(expected, xbox)
483
+ tm.assert_equal(lhs != NaT, expected)
484
+ tm.assert_equal(NaT != lhs, expected)
485
+
486
+ expected = np.array([False, False, False])
487
+ expected = tm.box_expected(expected, xbox)
488
+ tm.assert_equal(lhs < NaT, expected)
489
+ tm.assert_equal(NaT > lhs, expected)
490
+
491
+ def test_dti_cmp_nat_behaves_like_float_cmp_nan(self):
492
+ fidx1 = pd.Index([1.0, np.nan, 3.0, np.nan, 5.0, 7.0])
493
+ fidx2 = pd.Index([2.0, 3.0, np.nan, np.nan, 6.0, 7.0])
494
+
495
+ didx1 = DatetimeIndex(
496
+ ["2014-01-01", NaT, "2014-03-01", NaT, "2014-05-01", "2014-07-01"]
497
+ )
498
+ didx2 = DatetimeIndex(
499
+ ["2014-02-01", "2014-03-01", NaT, NaT, "2014-06-01", "2014-07-01"]
500
+ )
501
+ darr = np.array(
502
+ [
503
+ np.datetime64("2014-02-01 00:00"),
504
+ np.datetime64("2014-03-01 00:00"),
505
+ np.datetime64("nat", "ns"),
506
+ np.datetime64("nat", "ns"),
507
+ np.datetime64("2014-06-01 00:00"),
508
+ np.datetime64("2014-07-01 00:00"),
509
+ ]
510
+ )
511
+
512
+ cases = [(fidx1, fidx2), (didx1, didx2), (didx1, darr)]
513
+
514
+ # Check pd.NaT is handles as the same as np.nan
515
+ with tm.assert_produces_warning(None):
516
+ for idx1, idx2 in cases:
517
+ result = idx1 < idx2
518
+ expected = np.array([True, False, False, False, True, False])
519
+ tm.assert_numpy_array_equal(result, expected)
520
+
521
+ result = idx2 > idx1
522
+ expected = np.array([True, False, False, False, True, False])
523
+ tm.assert_numpy_array_equal(result, expected)
524
+
525
+ result = idx1 <= idx2
526
+ expected = np.array([True, False, False, False, True, True])
527
+ tm.assert_numpy_array_equal(result, expected)
528
+
529
+ result = idx2 >= idx1
530
+ expected = np.array([True, False, False, False, True, True])
531
+ tm.assert_numpy_array_equal(result, expected)
532
+
533
+ result = idx1 == idx2
534
+ expected = np.array([False, False, False, False, False, True])
535
+ tm.assert_numpy_array_equal(result, expected)
536
+
537
+ result = idx1 != idx2
538
+ expected = np.array([True, True, True, True, True, False])
539
+ tm.assert_numpy_array_equal(result, expected)
540
+
541
+ with tm.assert_produces_warning(None):
542
+ for idx1, val in [(fidx1, np.nan), (didx1, NaT)]:
543
+ result = idx1 < val
544
+ expected = np.array([False, False, False, False, False, False])
545
+ tm.assert_numpy_array_equal(result, expected)
546
+ result = idx1 > val
547
+ tm.assert_numpy_array_equal(result, expected)
548
+
549
+ result = idx1 <= val
550
+ tm.assert_numpy_array_equal(result, expected)
551
+ result = idx1 >= val
552
+ tm.assert_numpy_array_equal(result, expected)
553
+
554
+ result = idx1 == val
555
+ tm.assert_numpy_array_equal(result, expected)
556
+
557
+ result = idx1 != val
558
+ expected = np.array([True, True, True, True, True, True])
559
+ tm.assert_numpy_array_equal(result, expected)
560
+
561
+ # Check pd.NaT is handles as the same as np.nan
562
+ with tm.assert_produces_warning(None):
563
+ for idx1, val in [(fidx1, 3), (didx1, datetime(2014, 3, 1))]:
564
+ result = idx1 < val
565
+ expected = np.array([True, False, False, False, False, False])
566
+ tm.assert_numpy_array_equal(result, expected)
567
+ result = idx1 > val
568
+ expected = np.array([False, False, False, False, True, True])
569
+ tm.assert_numpy_array_equal(result, expected)
570
+
571
+ result = idx1 <= val
572
+ expected = np.array([True, False, True, False, False, False])
573
+ tm.assert_numpy_array_equal(result, expected)
574
+ result = idx1 >= val
575
+ expected = np.array([False, False, True, False, True, True])
576
+ tm.assert_numpy_array_equal(result, expected)
577
+
578
+ result = idx1 == val
579
+ expected = np.array([False, False, True, False, False, False])
580
+ tm.assert_numpy_array_equal(result, expected)
581
+
582
+ result = idx1 != val
583
+ expected = np.array([True, True, False, True, True, True])
584
+ tm.assert_numpy_array_equal(result, expected)
585
+
586
+ def test_comparison_tzawareness_compat(self, comparison_op, box_with_array):
587
+ # GH#18162
588
+ op = comparison_op
589
+ box = box_with_array
590
+
591
+ dr = date_range("2016-01-01", periods=6, unit="ns")
592
+ dz = dr.tz_localize("US/Pacific")
593
+
594
+ dr = tm.box_expected(dr, box)
595
+ dz = tm.box_expected(dz, box)
596
+
597
+ if box is pd.DataFrame:
598
+ tolist = lambda x: x.astype(object).values.tolist()[0]
599
+ else:
600
+ tolist = list
601
+
602
+ if op not in [operator.eq, operator.ne]:
603
+ msg = (
604
+ r"Invalid comparison between dtype=datetime64\[ns.*\] "
605
+ "and (Timestamp|DatetimeArray|list|ndarray)"
606
+ )
607
+ with pytest.raises(TypeError, match=msg):
608
+ op(dr, dz)
609
+
610
+ with pytest.raises(TypeError, match=msg):
611
+ op(dr, tolist(dz))
612
+ with pytest.raises(TypeError, match=msg):
613
+ op(dr, np.array(tolist(dz), dtype=object))
614
+ with pytest.raises(TypeError, match=msg):
615
+ op(dz, dr)
616
+
617
+ with pytest.raises(TypeError, match=msg):
618
+ op(dz, tolist(dr))
619
+ with pytest.raises(TypeError, match=msg):
620
+ op(dz, np.array(tolist(dr), dtype=object))
621
+
622
+ # The aware==aware and naive==naive comparisons should *not* raise
623
+ assert np.all(dr == dr)
624
+ assert np.all(dr == tolist(dr))
625
+ assert np.all(tolist(dr) == dr)
626
+ assert np.all(np.array(tolist(dr), dtype=object) == dr)
627
+ assert np.all(dr == np.array(tolist(dr), dtype=object))
628
+
629
+ assert np.all(dz == dz)
630
+ assert np.all(dz == tolist(dz))
631
+ assert np.all(tolist(dz) == dz)
632
+ assert np.all(np.array(tolist(dz), dtype=object) == dz)
633
+ assert np.all(dz == np.array(tolist(dz), dtype=object))
634
+
635
+ def test_comparison_tzawareness_compat_scalars(self, comparison_op, box_with_array):
636
+ # GH#18162
637
+ op = comparison_op
638
+
639
+ dr = date_range("2016-01-01", periods=6, unit="ns")
640
+ dz = dr.tz_localize("US/Pacific")
641
+
642
+ dr = tm.box_expected(dr, box_with_array)
643
+ dz = tm.box_expected(dz, box_with_array)
644
+
645
+ # Check comparisons against scalar Timestamps
646
+ ts = Timestamp("2000-03-14 01:59")
647
+ ts_tz = Timestamp("2000-03-14 01:59", tz="Europe/Amsterdam")
648
+
649
+ assert np.all(dr > ts)
650
+ msg = r"Invalid comparison between dtype=datetime64\[ns.*\] and Timestamp"
651
+ if op not in [operator.eq, operator.ne]:
652
+ with pytest.raises(TypeError, match=msg):
653
+ op(dr, ts_tz)
654
+
655
+ assert np.all(dz > ts_tz)
656
+ if op not in [operator.eq, operator.ne]:
657
+ with pytest.raises(TypeError, match=msg):
658
+ op(dz, ts)
659
+
660
+ if op not in [operator.eq, operator.ne]:
661
+ # GH#12601: Check comparison against Timestamps and DatetimeIndex
662
+ with pytest.raises(TypeError, match=msg):
663
+ op(ts, dz)
664
+
665
+ @pytest.mark.parametrize(
666
+ "other",
667
+ [datetime(2016, 1, 1), Timestamp("2016-01-01"), np.datetime64("2016-01-01")],
668
+ )
669
+ # Bug in NumPy? https://github.com/numpy/numpy/issues/13841
670
+ # Raising in __eq__ will fallback to NumPy, which warns, fails,
671
+ # then re-raises the original exception. So we just need to ignore.
672
+ @pytest.mark.filterwarnings("ignore:elementwise comp:DeprecationWarning")
673
+ def test_scalar_comparison_tzawareness(
674
+ self, comparison_op, other, tz_aware_fixture, box_with_array
675
+ ):
676
+ op = comparison_op
677
+ tz = tz_aware_fixture
678
+ dti = date_range("2016-01-01", periods=2, tz=tz, unit="ns")
679
+
680
+ dtarr = tm.box_expected(dti, box_with_array)
681
+ xbox = get_upcast_box(dtarr, other, True)
682
+ if op in [operator.eq, operator.ne]:
683
+ exbool = op is operator.ne
684
+ expected = np.array([exbool, exbool], dtype=bool)
685
+ expected = tm.box_expected(expected, xbox)
686
+
687
+ result = op(dtarr, other)
688
+ tm.assert_equal(result, expected)
689
+
690
+ result = op(other, dtarr)
691
+ tm.assert_equal(result, expected)
692
+ else:
693
+ msg = (
694
+ r"Invalid comparison between dtype=datetime64\[ns, .*\] "
695
+ f"and {type(other).__name__}"
696
+ )
697
+ with pytest.raises(TypeError, match=msg):
698
+ op(dtarr, other)
699
+ with pytest.raises(TypeError, match=msg):
700
+ op(other, dtarr)
701
+
702
+ def test_nat_comparison_tzawareness(self, comparison_op):
703
+ # GH#19276
704
+ # tzaware DatetimeIndex should not raise when compared to NaT
705
+ op = comparison_op
706
+
707
+ dti = DatetimeIndex(
708
+ ["2014-01-01", NaT, "2014-03-01", NaT, "2014-05-01", "2014-07-01"]
709
+ )
710
+ expected = np.array([op == operator.ne] * len(dti))
711
+ result = op(dti, NaT)
712
+ tm.assert_numpy_array_equal(result, expected)
713
+
714
+ result = op(dti.tz_localize("US/Pacific"), NaT)
715
+ tm.assert_numpy_array_equal(result, expected)
716
+
717
+ def test_dti_cmp_str(self, tz_naive_fixture):
718
+ # GH#22074
719
+ # regardless of tz, we expect these comparisons are valid
720
+ tz = tz_naive_fixture
721
+ rng = date_range("1/1/2000", periods=10, tz=tz)
722
+ other = "1/1/2000"
723
+
724
+ result = rng == other
725
+ expected = np.array([True] + [False] * 9)
726
+ tm.assert_numpy_array_equal(result, expected)
727
+
728
+ result = rng != other
729
+ expected = np.array([False] + [True] * 9)
730
+ tm.assert_numpy_array_equal(result, expected)
731
+
732
+ result = rng < other
733
+ expected = np.array([False] * 10)
734
+ tm.assert_numpy_array_equal(result, expected)
735
+
736
+ result = rng <= other
737
+ expected = np.array([True] + [False] * 9)
738
+ tm.assert_numpy_array_equal(result, expected)
739
+
740
+ result = rng > other
741
+ expected = np.array([False] + [True] * 9)
742
+ tm.assert_numpy_array_equal(result, expected)
743
+
744
+ result = rng >= other
745
+ expected = np.array([True] * 10)
746
+ tm.assert_numpy_array_equal(result, expected)
747
+
748
+ def test_dti_cmp_list(self):
749
+ rng = date_range("1/1/2000", periods=10)
750
+
751
+ result = rng == list(rng)
752
+ expected = rng == rng
753
+ tm.assert_numpy_array_equal(result, expected)
754
+
755
+ @pytest.mark.parametrize(
756
+ "other",
757
+ [
758
+ pd.timedelta_range("1D", periods=10),
759
+ pd.timedelta_range("1D", periods=10).to_series(),
760
+ pd.timedelta_range("1D", periods=10).asi8.view("m8[ns]"),
761
+ ],
762
+ ids=lambda x: type(x).__name__,
763
+ )
764
+ def test_dti_cmp_tdi_tzawareness(self, other):
765
+ # GH#22074
766
+ # reversion test that we _don't_ call _assert_tzawareness_compat
767
+ # when comparing against TimedeltaIndex
768
+ dti = date_range("2000-01-01", periods=10, tz="Asia/Tokyo")
769
+
770
+ result = dti == other
771
+ expected = np.array([False] * 10)
772
+ if isinstance(other, Series):
773
+ tm.assert_series_equal(result, Series(expected, index=other.index))
774
+ else:
775
+ tm.assert_numpy_array_equal(result, expected)
776
+
777
+ result = dti != other
778
+ expected = np.array([True] * 10)
779
+ if isinstance(other, Series):
780
+ tm.assert_series_equal(result, Series(expected, index=other.index))
781
+ else:
782
+ tm.assert_numpy_array_equal(result, expected)
783
+
784
+ msg = "Invalid comparison between"
785
+ with pytest.raises(TypeError, match=msg):
786
+ dti < other
787
+ with pytest.raises(TypeError, match=msg):
788
+ dti <= other
789
+ with pytest.raises(TypeError, match=msg):
790
+ dti > other
791
+ with pytest.raises(TypeError, match=msg):
792
+ dti >= other
793
+
794
+ def test_dti_cmp_object_dtype(self):
795
+ # GH#22074
796
+ dti = date_range("2000-01-01", periods=10, tz="Asia/Tokyo")
797
+
798
+ other = dti.astype("O")
799
+
800
+ result = dti == other
801
+ expected = np.array([True] * 10)
802
+ tm.assert_numpy_array_equal(result, expected)
803
+
804
+ other = dti.tz_localize(None)
805
+ result = dti != other
806
+ tm.assert_numpy_array_equal(result, expected)
807
+
808
+ other = np.array(list(dti[:5]) + [Timedelta(days=1)] * 5)
809
+ result = dti == other
810
+ expected = np.array([True] * 5 + [False] * 5)
811
+ tm.assert_numpy_array_equal(result, expected)
812
+ msg = ">=' not supported between instances of 'Timestamp' and 'Timedelta'"
813
+ with pytest.raises(TypeError, match=msg):
814
+ dti >= other
815
+
816
+
817
+ # ------------------------------------------------------------------
818
+ # Arithmetic
819
+
820
+
821
+ class TestDatetime64Arithmetic:
822
+ # This class is intended for "finished" tests that are fully parametrized
823
+ # over DataFrame/Series/Index/DatetimeArray
824
+
825
+ # -------------------------------------------------------------
826
+ # Addition/Subtraction of timedelta-like
827
+
828
+ @pytest.mark.arm_slow
829
+ def test_dt64arr_add_timedeltalike_scalar(
830
+ self, tz_naive_fixture, two_hours, box_with_array
831
+ ):
832
+ # GH#22005, GH#22163 check DataFrame doesn't raise TypeError
833
+ tz = tz_naive_fixture
834
+
835
+ rng = date_range("2000-01-01", "2000-02-01", tz=tz, unit="ns")
836
+ expected = date_range("2000-01-01 02:00", "2000-02-01 02:00", tz=tz, unit="ns")
837
+ if tz is not None:
838
+ expected = expected._with_freq(None)
839
+
840
+ rng = tm.box_expected(rng, box_with_array)
841
+ expected = tm.box_expected(expected, box_with_array)
842
+
843
+ result = rng + two_hours
844
+ tm.assert_equal(result, expected)
845
+
846
+ result = two_hours + rng
847
+ tm.assert_equal(result, expected)
848
+
849
+ rng += two_hours
850
+ tm.assert_equal(rng, expected)
851
+
852
+ def test_dt64arr_sub_timedeltalike_scalar(
853
+ self, tz_naive_fixture, two_hours, box_with_array
854
+ ):
855
+ tz = tz_naive_fixture
856
+
857
+ rng = date_range("2000-01-01", "2000-02-01", tz=tz, unit="ns")
858
+ expected = date_range("1999-12-31 22:00", "2000-01-31 22:00", tz=tz, unit="ns")
859
+ if tz is not None:
860
+ expected = expected._with_freq(None)
861
+
862
+ rng = tm.box_expected(rng, box_with_array)
863
+ expected = tm.box_expected(expected, box_with_array)
864
+
865
+ result = rng - two_hours
866
+ tm.assert_equal(result, expected)
867
+
868
+ rng -= two_hours
869
+ tm.assert_equal(rng, expected)
870
+
871
+ def test_dt64_array_sub_dt_with_different_timezone(self, box_with_array):
872
+ t1 = date_range("20130101", periods=3).tz_localize("US/Eastern")
873
+ t1 = tm.box_expected(t1, box_with_array)
874
+ t2 = Timestamp("20130101").tz_localize("CET")
875
+ tnaive = Timestamp(20130101)
876
+
877
+ result = t1 - t2
878
+ expected = TimedeltaIndex(
879
+ ["0 days 06:00:00", "1 days 06:00:00", "2 days 06:00:00"]
880
+ )
881
+ expected = tm.box_expected(expected, box_with_array)
882
+ tm.assert_equal(result, expected)
883
+
884
+ result = t2 - t1
885
+ expected = TimedeltaIndex(
886
+ ["-1 days +18:00:00", "-2 days +18:00:00", "-3 days +18:00:00"]
887
+ )
888
+ expected = tm.box_expected(expected, box_with_array)
889
+ tm.assert_equal(result, expected)
890
+
891
+ msg = "Cannot subtract tz-naive and tz-aware datetime-like objects"
892
+ with pytest.raises(TypeError, match=msg):
893
+ t1 - tnaive
894
+
895
+ with pytest.raises(TypeError, match=msg):
896
+ tnaive - t1
897
+
898
+ def test_dt64_array_sub_dt64_array_with_different_timezone(self, box_with_array):
899
+ t1 = date_range("20130101", periods=3).tz_localize("US/Eastern")
900
+ t1 = tm.box_expected(t1, box_with_array)
901
+ t2 = date_range("20130101", periods=3).tz_localize("CET")
902
+ t2 = tm.box_expected(t2, box_with_array)
903
+ tnaive = date_range("20130101", periods=3)
904
+
905
+ result = t1 - t2
906
+ expected = TimedeltaIndex(
907
+ ["0 days 06:00:00", "0 days 06:00:00", "0 days 06:00:00"]
908
+ )
909
+ expected = tm.box_expected(expected, box_with_array)
910
+ tm.assert_equal(result, expected)
911
+
912
+ result = t2 - t1
913
+ expected = TimedeltaIndex(
914
+ ["-1 days +18:00:00", "-1 days +18:00:00", "-1 days +18:00:00"]
915
+ )
916
+ expected = tm.box_expected(expected, box_with_array)
917
+ tm.assert_equal(result, expected)
918
+
919
+ msg = "Cannot subtract tz-naive and tz-aware datetime-like objects"
920
+ with pytest.raises(TypeError, match=msg):
921
+ t1 - tnaive
922
+
923
+ with pytest.raises(TypeError, match=msg):
924
+ tnaive - t1
925
+
926
+ def test_dt64arr_add_sub_td64_nat(self, box_with_array, tz_naive_fixture):
927
+ # GH#23320 special handling for timedelta64("NaT")
928
+ tz = tz_naive_fixture
929
+
930
+ dti = date_range("1994-04-01", periods=9, tz=tz, freq="QS", unit="ns")
931
+ other = np.timedelta64("NaT", "ns")
932
+ expected = DatetimeIndex(["NaT"] * 9, tz=tz).as_unit("ns")
933
+
934
+ obj = tm.box_expected(dti, box_with_array)
935
+ expected = tm.box_expected(expected, box_with_array)
936
+
937
+ result = obj + other
938
+ tm.assert_equal(result, expected)
939
+ result = other + obj
940
+ tm.assert_equal(result, expected)
941
+ result = obj - other
942
+ tm.assert_equal(result, expected)
943
+ msg = "cannot subtract"
944
+ with pytest.raises(TypeError, match=msg):
945
+ other - obj
946
+
947
+ def test_dt64arr_add_sub_td64ndarray(self, tz_naive_fixture, box_with_array):
948
+ tz = tz_naive_fixture
949
+ dti = date_range("2016-01-01", periods=3, tz=tz)
950
+ tdi = TimedeltaIndex(["-1 Day", "-1 Day", "-1 Day"])
951
+ tdarr = tdi.values
952
+
953
+ expected = date_range("2015-12-31", "2016-01-02", periods=3, tz=tz)
954
+
955
+ dtarr = tm.box_expected(dti, box_with_array)
956
+ expected = tm.box_expected(expected, box_with_array)
957
+
958
+ result = dtarr + tdarr
959
+ tm.assert_equal(result, expected)
960
+ result = tdarr + dtarr
961
+ tm.assert_equal(result, expected)
962
+
963
+ expected = date_range("2016-01-02", "2016-01-04", periods=3, tz=tz)
964
+ expected = tm.box_expected(expected, box_with_array)
965
+
966
+ result = dtarr - tdarr
967
+ tm.assert_equal(result, expected)
968
+ msg = "|".join(
969
+ [
970
+ "cannot subtract DatetimeArray from ndarray",
971
+ "cannot subtract a datelike from a TimedeltaArray",
972
+ "cannot subtract DatetimeArray from Timedelta",
973
+ ]
974
+ )
975
+ with pytest.raises(TypeError, match=msg):
976
+ tdarr - dtarr
977
+
978
+ # -----------------------------------------------------------------
979
+ # Subtraction of datetime-like scalars
980
+
981
+ @pytest.mark.parametrize(
982
+ "ts",
983
+ [
984
+ Timestamp("2013-01-01"),
985
+ Timestamp("2013-01-01").to_pydatetime(),
986
+ Timestamp("2013-01-01").to_datetime64(),
987
+ # GH#7996, GH#22163 ensure non-nano datetime64 is converted to nano
988
+ # for DataFrame operation
989
+ np.datetime64("2013-01-01", "D"),
990
+ ],
991
+ )
992
+ def test_dt64arr_sub_dtscalar(self, box_with_array, ts):
993
+ # GH#8554, GH#22163 DataFrame op should _not_ return dt64 dtype
994
+ idx = date_range("2013-01-01", periods=3)._with_freq(None)
995
+ idx = tm.box_expected(idx, box_with_array)
996
+
997
+ expected = TimedeltaIndex(["0 Days", "1 Day", "2 Days"])
998
+ expected = tm.box_expected(expected, box_with_array)
999
+
1000
+ result = idx - ts
1001
+ tm.assert_equal(result, expected)
1002
+
1003
+ result = ts - idx
1004
+ tm.assert_equal(result, -expected)
1005
+ tm.assert_equal(result, -expected)
1006
+
1007
+ def test_dt64arr_sub_timestamp_tzaware(self, box_with_array):
1008
+ ser = date_range("2014-03-17", periods=2, freq="D", tz="US/Eastern", unit="ns")
1009
+ ser = ser._with_freq(None)
1010
+ ts = ser[0]
1011
+
1012
+ ser = tm.box_expected(ser, box_with_array)
1013
+
1014
+ delta_series = Series(
1015
+ [np.timedelta64(0, "D"), np.timedelta64(1, "D")], dtype="m8[ns]"
1016
+ )
1017
+ expected = tm.box_expected(delta_series, box_with_array)
1018
+
1019
+ tm.assert_equal(ser - ts, expected)
1020
+ tm.assert_equal(ts - ser, -expected)
1021
+
1022
+ def test_dt64arr_sub_NaT(self, box_with_array, unit):
1023
+ # GH#18808
1024
+ dti = DatetimeIndex([NaT, Timestamp("19900315")]).as_unit(unit)
1025
+ ser = tm.box_expected(dti, box_with_array)
1026
+
1027
+ result = ser - NaT
1028
+ expected = Series([NaT, NaT], dtype=f"timedelta64[{unit}]")
1029
+ expected = tm.box_expected(expected, box_with_array)
1030
+ tm.assert_equal(result, expected)
1031
+
1032
+ dti_tz = dti.tz_localize("Asia/Tokyo")
1033
+ ser_tz = tm.box_expected(dti_tz, box_with_array)
1034
+
1035
+ result = ser_tz - NaT
1036
+ expected = Series([NaT, NaT], dtype=f"timedelta64[{unit}]")
1037
+ expected = tm.box_expected(expected, box_with_array)
1038
+ tm.assert_equal(result, expected)
1039
+
1040
+ # -------------------------------------------------------------
1041
+ # Subtraction of datetime-like array-like
1042
+
1043
+ def test_dt64arr_sub_dt64object_array(
1044
+ self, performance_warning, box_with_array, tz_naive_fixture
1045
+ ):
1046
+ dti = date_range("2016-01-01", periods=3, tz=tz_naive_fixture)
1047
+ expected = dti - dti
1048
+
1049
+ obj = tm.box_expected(dti, box_with_array)
1050
+ expected = tm.box_expected(expected, box_with_array).astype(object)
1051
+
1052
+ with tm.assert_produces_warning(performance_warning):
1053
+ result = obj - obj.astype(object)
1054
+ tm.assert_equal(result, expected)
1055
+
1056
+ def test_dt64arr_naive_sub_dt64ndarray(self, box_with_array):
1057
+ dti = date_range("2016-01-01", periods=3, tz=None)
1058
+ dt64vals = dti.values
1059
+
1060
+ dtarr = tm.box_expected(dti, box_with_array)
1061
+
1062
+ expected = dtarr - dtarr
1063
+ result = dtarr - dt64vals
1064
+ tm.assert_equal(result, expected)
1065
+ result = dt64vals - dtarr
1066
+ tm.assert_equal(result, expected)
1067
+
1068
+ def test_dt64arr_aware_sub_dt64ndarray_raises(
1069
+ self, tz_aware_fixture, box_with_array
1070
+ ):
1071
+ tz = tz_aware_fixture
1072
+ dti = date_range("2016-01-01", periods=3, tz=tz)
1073
+ dt64vals = dti.values
1074
+
1075
+ dtarr = tm.box_expected(dti, box_with_array)
1076
+ msg = "Cannot subtract tz-naive and tz-aware datetime"
1077
+ with pytest.raises(TypeError, match=msg):
1078
+ dtarr - dt64vals
1079
+ with pytest.raises(TypeError, match=msg):
1080
+ dt64vals - dtarr
1081
+
1082
+ # -------------------------------------------------------------
1083
+ # Addition of datetime-like others (invalid)
1084
+
1085
+ def test_dt64arr_add_dtlike_raises(self, tz_naive_fixture, box_with_array):
1086
+ # GH#22163 ensure DataFrame doesn't cast Timestamp to i8
1087
+ # GH#9631
1088
+ tz = tz_naive_fixture
1089
+
1090
+ dti = date_range("2016-01-01", periods=3, tz=tz)
1091
+ if tz is None:
1092
+ dti2 = dti.tz_localize("US/Eastern")
1093
+ else:
1094
+ dti2 = dti.tz_localize(None)
1095
+ dtarr = tm.box_expected(dti, box_with_array)
1096
+
1097
+ assert_cannot_add(dtarr, dti.values)
1098
+ assert_cannot_add(dtarr, dti)
1099
+ assert_cannot_add(dtarr, dtarr)
1100
+ assert_cannot_add(dtarr, dti[0])
1101
+ assert_cannot_add(dtarr, dti[0].to_pydatetime())
1102
+ assert_cannot_add(dtarr, dti[0].to_datetime64())
1103
+ assert_cannot_add(dtarr, dti2[0])
1104
+ assert_cannot_add(dtarr, dti2[0].to_pydatetime())
1105
+ assert_cannot_add(dtarr, np.datetime64("2011-01-01", "D"))
1106
+
1107
+ # -------------------------------------------------------------
1108
+ # Other Invalid Addition/Subtraction
1109
+
1110
+ # Note: freq here includes both Tick and non-Tick offsets; this is
1111
+ # relevant because historically integer-addition was allowed if we had
1112
+ # a freq.
1113
+ @pytest.mark.parametrize("freq", ["h", "D", "W", "2ME", "MS", "QE", "B", None])
1114
+ @pytest.mark.parametrize("dtype", [None, "uint8"])
1115
+ def test_dt64arr_addsub_intlike(
1116
+ self, dtype, index_or_series_or_array, freq, tz_naive_fixture
1117
+ ):
1118
+ # GH#19959, GH#19123, GH#19012
1119
+ # GH#55860 use index_or_series_or_array instead of box_with_array
1120
+ # bc DataFrame alignment makes it inapplicable
1121
+ tz = tz_naive_fixture
1122
+
1123
+ if freq is None:
1124
+ dti = DatetimeIndex(["NaT", "2017-04-05 06:07:08"], tz=tz)
1125
+ else:
1126
+ dti = date_range("2016-01-01", periods=2, freq=freq, tz=tz)
1127
+
1128
+ obj = index_or_series_or_array(dti)
1129
+ other = np.array([4, -1])
1130
+ if dtype is not None:
1131
+ other = other.astype(dtype)
1132
+
1133
+ msg = "|".join(
1134
+ [
1135
+ "Addition/subtraction of integers",
1136
+ "cannot subtract DatetimeArray from",
1137
+ # IntegerArray
1138
+ "can only perform ops with numeric values",
1139
+ "unsupported operand type.*Categorical",
1140
+ r"unsupported operand type\(s\) for -: 'int' and 'Timestamp'",
1141
+ ]
1142
+ )
1143
+ assert_invalid_addsub_type(obj, 1, msg)
1144
+ assert_invalid_addsub_type(obj, np.int64(2), msg)
1145
+ assert_invalid_addsub_type(obj, np.array(3, dtype=np.int64), msg)
1146
+ assert_invalid_addsub_type(obj, other, msg)
1147
+ assert_invalid_addsub_type(obj, np.array(other), msg)
1148
+ assert_invalid_addsub_type(obj, pd.array(other), msg)
1149
+ assert_invalid_addsub_type(obj, pd.Categorical(other), msg)
1150
+ assert_invalid_addsub_type(obj, pd.Index(other), msg)
1151
+ assert_invalid_addsub_type(obj, Series(other), msg)
1152
+
1153
+ @pytest.mark.parametrize(
1154
+ "other",
1155
+ [
1156
+ 3.14,
1157
+ np.array([2.0, 3.0]),
1158
+ # GH#13078 datetime +/- Period is invalid
1159
+ Period("2011-01-01", freq="D"),
1160
+ # https://github.com/pandas-dev/pandas/issues/10329
1161
+ time(1, 2, 3),
1162
+ ],
1163
+ )
1164
+ @pytest.mark.parametrize("dti_freq", [None, "D"])
1165
+ def test_dt64arr_add_sub_invalid(self, dti_freq, other, box_with_array):
1166
+ dti = DatetimeIndex(["2011-01-01", "2011-01-02"], freq=dti_freq)
1167
+ dtarr = tm.box_expected(dti, box_with_array)
1168
+ msg = "|".join(
1169
+ [
1170
+ "unsupported operand type",
1171
+ "cannot (add|subtract)",
1172
+ "cannot use operands with types",
1173
+ "ufunc '?(add|subtract)'? cannot use operands with types",
1174
+ "Concatenation operation is not implemented for NumPy arrays",
1175
+ ]
1176
+ )
1177
+ assert_invalid_addsub_type(dtarr, other, msg)
1178
+
1179
+ @pytest.mark.parametrize("pi_freq", ["D", "W", "Q", "h"])
1180
+ @pytest.mark.parametrize("dti_freq", [None, "D"])
1181
+ def test_dt64arr_add_sub_parr(
1182
+ self, dti_freq, pi_freq, box_with_array, box_with_array2
1183
+ ):
1184
+ # GH#20049 subtracting PeriodIndex should raise TypeError
1185
+ dti = DatetimeIndex(["2011-01-01", "2011-01-02"], freq=dti_freq)
1186
+ pi = dti.to_period(pi_freq)
1187
+
1188
+ dtarr = tm.box_expected(dti, box_with_array)
1189
+ parr = tm.box_expected(pi, box_with_array2)
1190
+ msg = "|".join(
1191
+ [
1192
+ "cannot (add|subtract)",
1193
+ "unsupported operand",
1194
+ "descriptor.*requires",
1195
+ "ufunc.*cannot use operands",
1196
+ ]
1197
+ )
1198
+ assert_invalid_addsub_type(dtarr, parr, msg)
1199
+
1200
+ @pytest.mark.filterwarnings("ignore::pandas.errors.PerformanceWarning")
1201
+ def test_dt64arr_addsub_time_objects_raises(self, box_with_array, tz_naive_fixture):
1202
+ # https://github.com/pandas-dev/pandas/issues/10329
1203
+
1204
+ tz = tz_naive_fixture
1205
+
1206
+ obj1 = date_range("2012-01-01", periods=3, tz=tz)
1207
+ obj2 = [time(i, i, i) for i in range(3)]
1208
+
1209
+ obj1 = tm.box_expected(obj1, box_with_array)
1210
+ obj2 = tm.box_expected(obj2, box_with_array)
1211
+
1212
+ msg = "|".join(
1213
+ [
1214
+ "unsupported operand",
1215
+ "cannot subtract DatetimeArray from ndarray",
1216
+ ]
1217
+ )
1218
+ # pandas.errors.PerformanceWarning: Non-vectorized DateOffset being
1219
+ # applied to Series or DatetimeIndex
1220
+ # we aren't testing that here, so ignore.
1221
+ assert_invalid_addsub_type(obj1, obj2, msg=msg)
1222
+
1223
+ # -------------------------------------------------------------
1224
+ # Other invalid operations
1225
+
1226
+ @pytest.mark.parametrize(
1227
+ "dt64_series",
1228
+ [
1229
+ Series([Timestamp("19900315"), Timestamp("19900315")]),
1230
+ Series([NaT, Timestamp("19900315")]),
1231
+ Series([NaT, NaT], dtype="datetime64[ns]"),
1232
+ ],
1233
+ )
1234
+ @pytest.mark.parametrize("one", [1, 1.0, np.array(1)])
1235
+ def test_dt64_mul_div_numeric_invalid(self, one, dt64_series, box_with_array):
1236
+ obj = tm.box_expected(dt64_series, box_with_array)
1237
+
1238
+ msg = "cannot perform .* with this index type"
1239
+
1240
+ # multiplication
1241
+ with pytest.raises(TypeError, match=msg):
1242
+ obj * one
1243
+ with pytest.raises(TypeError, match=msg):
1244
+ one * obj
1245
+
1246
+ # division
1247
+ with pytest.raises(TypeError, match=msg):
1248
+ obj / one
1249
+ with pytest.raises(TypeError, match=msg):
1250
+ one / obj
1251
+
1252
+
1253
+ class TestDatetime64DateOffsetArithmetic:
1254
+ # -------------------------------------------------------------
1255
+ # Tick DateOffsets
1256
+
1257
+ # TODO: parametrize over timezone?
1258
+ def test_dt64arr_series_add_tick_DateOffset(self, box_with_array, unit):
1259
+ # GH#4532
1260
+ # operate with pd.offsets
1261
+ ser = Series(
1262
+ [Timestamp("20130101 9:01"), Timestamp("20130101 9:02")]
1263
+ ).dt.as_unit(unit)
1264
+ expected = Series(
1265
+ [Timestamp("20130101 9:01:05"), Timestamp("20130101 9:02:05")]
1266
+ ).dt.as_unit(unit)
1267
+
1268
+ ser = tm.box_expected(ser, box_with_array)
1269
+ expected = tm.box_expected(expected, box_with_array)
1270
+
1271
+ result = ser + pd.offsets.Second(5)
1272
+ tm.assert_equal(result, expected)
1273
+
1274
+ result2 = pd.offsets.Second(5) + ser
1275
+ tm.assert_equal(result2, expected)
1276
+
1277
+ def test_dt64arr_series_sub_tick_DateOffset(self, box_with_array):
1278
+ # GH#4532
1279
+ # operate with pd.offsets
1280
+ ser = Series([Timestamp("20130101 9:01"), Timestamp("20130101 9:02")])
1281
+ expected = Series(
1282
+ [Timestamp("20130101 9:00:55"), Timestamp("20130101 9:01:55")]
1283
+ )
1284
+
1285
+ ser = tm.box_expected(ser, box_with_array)
1286
+ expected = tm.box_expected(expected, box_with_array)
1287
+
1288
+ result = ser - pd.offsets.Second(5)
1289
+ tm.assert_equal(result, expected)
1290
+
1291
+ result2 = -pd.offsets.Second(5) + ser
1292
+ tm.assert_equal(result2, expected)
1293
+ msg = "cannot subtract DatetimeArray from Second"
1294
+ with pytest.raises(TypeError, match=msg):
1295
+ pd.offsets.Second(5) - ser
1296
+
1297
+ @pytest.mark.parametrize(
1298
+ "cls_name", ["Day", "Hour", "Minute", "Second", "Milli", "Micro", "Nano"]
1299
+ )
1300
+ def test_dt64arr_add_sub_tick_DateOffset_smoke(self, cls_name, box_with_array):
1301
+ # GH#4532
1302
+ # smoke tests for valid DateOffsets
1303
+ ser = Series([Timestamp("20130101 9:01"), Timestamp("20130101 9:02")])
1304
+ ser = tm.box_expected(ser, box_with_array)
1305
+
1306
+ offset_cls = getattr(pd.offsets, cls_name)
1307
+ ser + offset_cls(5)
1308
+ offset_cls(5) + ser
1309
+ ser - offset_cls(5)
1310
+
1311
+ def test_dti_add_tick_tzaware(self, tz_aware_fixture, box_with_array):
1312
+ # GH#21610, GH#22163 ensure DataFrame doesn't return object-dtype
1313
+ tz = tz_aware_fixture
1314
+ if tz == "US/Pacific":
1315
+ dates = date_range("2012-11-01", periods=3, tz=tz, unit="ns")
1316
+ offset = dates + pd.offsets.Hour(5)
1317
+ assert dates[0] + pd.offsets.Hour(5) == offset[0]
1318
+
1319
+ dates = date_range("2010-11-01 00:00", periods=3, tz=tz, freq="h", unit="ns")
1320
+ expected = DatetimeIndex(
1321
+ ["2010-11-01 05:00", "2010-11-01 06:00", "2010-11-01 07:00"],
1322
+ freq="h",
1323
+ tz=tz,
1324
+ ).as_unit("ns")
1325
+
1326
+ dates = tm.box_expected(dates, box_with_array)
1327
+ expected = tm.box_expected(expected, box_with_array)
1328
+
1329
+ for scalar in [pd.offsets.Hour(5), np.timedelta64(5, "h"), timedelta(hours=5)]:
1330
+ offset = dates + scalar
1331
+ tm.assert_equal(offset, expected)
1332
+ offset = scalar + dates
1333
+ tm.assert_equal(offset, expected)
1334
+
1335
+ roundtrip = offset - scalar
1336
+ tm.assert_equal(roundtrip, dates)
1337
+
1338
+ msg = "cannot subtract DatetimeArray from"
1339
+ with pytest.raises(TypeError, match=msg):
1340
+ scalar - dates
1341
+
1342
+ # -------------------------------------------------------------
1343
+ # RelativeDelta DateOffsets
1344
+
1345
+ def test_dt64arr_add_sub_relativedelta_offsets(self, box_with_array, unit):
1346
+ # GH#10699
1347
+ vec = DatetimeIndex(
1348
+ [
1349
+ Timestamp("2000-01-05 00:15:00"),
1350
+ Timestamp("2000-01-31 00:23:00"),
1351
+ Timestamp("2000-01-01"),
1352
+ Timestamp("2000-03-31"),
1353
+ Timestamp("2000-02-29"),
1354
+ Timestamp("2000-12-31"),
1355
+ Timestamp("2000-05-15"),
1356
+ Timestamp("2001-06-15"),
1357
+ ]
1358
+ ).as_unit(unit)
1359
+ vec = tm.box_expected(vec, box_with_array)
1360
+ vec_items = vec.iloc[0] if box_with_array is pd.DataFrame else vec
1361
+
1362
+ # DateOffset relativedelta fastpath
1363
+ relative_kwargs = [
1364
+ ("years", 2),
1365
+ ("months", 5),
1366
+ ("days", 3),
1367
+ ("hours", 5),
1368
+ ("minutes", 10),
1369
+ ("seconds", 2),
1370
+ ("microseconds", 5),
1371
+ ]
1372
+ for i, (offset_unit, value) in enumerate(relative_kwargs):
1373
+ off = DateOffset(**{offset_unit: value})
1374
+
1375
+ exp_unit = unit
1376
+ if offset_unit == "microseconds" and unit != "ns":
1377
+ exp_unit = "us"
1378
+
1379
+ # TODO(GH#55564): as_unit will be unnecessary
1380
+ expected = DatetimeIndex([x + off for x in vec_items]).as_unit(exp_unit)
1381
+ expected = tm.box_expected(expected, box_with_array)
1382
+ tm.assert_equal(expected, vec + off)
1383
+
1384
+ expected = DatetimeIndex([x - off for x in vec_items]).as_unit(exp_unit)
1385
+ expected = tm.box_expected(expected, box_with_array)
1386
+ tm.assert_equal(expected, vec - off)
1387
+
1388
+ off = DateOffset(**dict(relative_kwargs[: i + 1]))
1389
+
1390
+ expected = DatetimeIndex([x + off for x in vec_items]).as_unit(exp_unit)
1391
+ expected = tm.box_expected(expected, box_with_array)
1392
+ tm.assert_equal(expected, vec + off)
1393
+
1394
+ expected = DatetimeIndex([x - off for x in vec_items]).as_unit(exp_unit)
1395
+ expected = tm.box_expected(expected, box_with_array)
1396
+ tm.assert_equal(expected, vec - off)
1397
+ msg = "cannot subtract DatetimeArray from"
1398
+ with pytest.raises(TypeError, match=msg):
1399
+ off - vec
1400
+
1401
+ # -------------------------------------------------------------
1402
+ # Non-Tick, Non-RelativeDelta DateOffsets
1403
+
1404
+ # TODO: redundant with test_dt64arr_add_sub_DateOffset? that includes
1405
+ # tz-aware cases which this does not
1406
+ @pytest.mark.filterwarnings("ignore::pandas.errors.PerformanceWarning")
1407
+ @pytest.mark.parametrize(
1408
+ "cls_and_kwargs",
1409
+ [
1410
+ "YearBegin",
1411
+ ("YearBegin", {"month": 5}),
1412
+ "YearEnd",
1413
+ ("YearEnd", {"month": 5}),
1414
+ "MonthBegin",
1415
+ "MonthEnd",
1416
+ "SemiMonthEnd",
1417
+ "SemiMonthBegin",
1418
+ "Week",
1419
+ ("Week", {"weekday": 3}),
1420
+ ("Week", {"weekday": 6}),
1421
+ "BusinessDay",
1422
+ "BDay",
1423
+ "QuarterEnd",
1424
+ "QuarterBegin",
1425
+ "CustomBusinessDay",
1426
+ "CDay",
1427
+ "CBMonthEnd",
1428
+ "CBMonthBegin",
1429
+ "BMonthBegin",
1430
+ "BMonthEnd",
1431
+ "BusinessHour",
1432
+ "BYearBegin",
1433
+ "BYearEnd",
1434
+ "BQuarterBegin",
1435
+ ("LastWeekOfMonth", {"weekday": 2}),
1436
+ (
1437
+ "FY5253Quarter",
1438
+ {
1439
+ "qtr_with_extra_week": 1,
1440
+ "startingMonth": 1,
1441
+ "weekday": 2,
1442
+ "variation": "nearest",
1443
+ },
1444
+ ),
1445
+ ("FY5253", {"weekday": 0, "startingMonth": 2, "variation": "nearest"}),
1446
+ ("WeekOfMonth", {"weekday": 2, "week": 2}),
1447
+ "Easter",
1448
+ ("DateOffset", {"day": 4}),
1449
+ ("DateOffset", {"month": 5}),
1450
+ ],
1451
+ )
1452
+ @pytest.mark.parametrize("normalize", [True, False])
1453
+ @pytest.mark.parametrize("n", [0, 5])
1454
+ @pytest.mark.parametrize("tz", [None, "US/Central"])
1455
+ def test_dt64arr_add_sub_DateOffsets(
1456
+ self, box_with_array, n, normalize, cls_and_kwargs, unit, tz
1457
+ ):
1458
+ # GH#10699
1459
+ # assert vectorized operation matches pointwise operations
1460
+
1461
+ if isinstance(cls_and_kwargs, tuple):
1462
+ # If cls_name param is a tuple, then 2nd entry is kwargs for
1463
+ # the offset constructor
1464
+ cls_name, kwargs = cls_and_kwargs
1465
+ else:
1466
+ cls_name = cls_and_kwargs
1467
+ kwargs = {}
1468
+
1469
+ if n == 0 and cls_name in [
1470
+ "WeekOfMonth",
1471
+ "LastWeekOfMonth",
1472
+ "FY5253Quarter",
1473
+ "FY5253",
1474
+ ]:
1475
+ # passing n = 0 is invalid for these offset classes
1476
+ return
1477
+
1478
+ vec = (
1479
+ DatetimeIndex(
1480
+ [
1481
+ Timestamp("2000-01-05 00:15:00"),
1482
+ Timestamp("2000-01-31 00:23:00"),
1483
+ Timestamp("2000-01-01"),
1484
+ Timestamp("2000-03-31"),
1485
+ Timestamp("2000-02-29"),
1486
+ Timestamp("2000-12-31"),
1487
+ Timestamp("2000-05-15"),
1488
+ Timestamp("2001-06-15"),
1489
+ ]
1490
+ )
1491
+ .as_unit(unit)
1492
+ .tz_localize(tz)
1493
+ )
1494
+ vec = tm.box_expected(vec, box_with_array)
1495
+ vec_items = vec.iloc[0] if box_with_array is pd.DataFrame else vec
1496
+
1497
+ offset_cls = getattr(pd.offsets, cls_name)
1498
+ offset = offset_cls(n, normalize=normalize, **kwargs)
1499
+
1500
+ # TODO(GH#55564): as_unit will be unnecessary
1501
+ expected = DatetimeIndex([x + offset for x in vec_items]).as_unit(unit)
1502
+ expected = tm.box_expected(expected, box_with_array)
1503
+ tm.assert_equal(expected, vec + offset)
1504
+ tm.assert_equal(expected, offset + vec)
1505
+
1506
+ expected = DatetimeIndex([x - offset for x in vec_items]).as_unit(unit)
1507
+ expected = tm.box_expected(expected, box_with_array)
1508
+ tm.assert_equal(expected, vec - offset)
1509
+
1510
+ expected = DatetimeIndex([offset + x for x in vec_items]).as_unit(unit)
1511
+ expected = tm.box_expected(expected, box_with_array)
1512
+ tm.assert_equal(expected, offset + vec)
1513
+ msg = "cannot subtract DatetimeArray from"
1514
+ with pytest.raises(TypeError, match=msg):
1515
+ offset - vec
1516
+
1517
+ @pytest.mark.parametrize(
1518
+ "other",
1519
+ [
1520
+ [pd.offsets.MonthEnd(), pd.offsets.Day(n=2)],
1521
+ [pd.offsets.DateOffset(years=1), pd.offsets.MonthEnd()],
1522
+ # matching offsets
1523
+ [pd.offsets.DateOffset(years=1), pd.offsets.DateOffset(years=1)],
1524
+ ],
1525
+ )
1526
+ @pytest.mark.parametrize("op", [operator.add, roperator.radd, operator.sub])
1527
+ def test_dt64arr_add_sub_offset_array(
1528
+ self, performance_warning, tz_naive_fixture, box_with_array, op, other
1529
+ ):
1530
+ # GH#18849
1531
+ # GH#10699 array of offsets
1532
+
1533
+ tz = tz_naive_fixture
1534
+ dti = date_range("2017-01-01", periods=2, tz=tz)
1535
+ dtarr = tm.box_expected(dti, box_with_array)
1536
+ other = np.array(other)
1537
+ expected = DatetimeIndex([op(dti[n], other[n]) for n in range(len(dti))])
1538
+ expected = tm.box_expected(expected, box_with_array).astype(object)
1539
+
1540
+ with tm.assert_produces_warning(performance_warning):
1541
+ res = op(dtarr, other)
1542
+ tm.assert_equal(res, expected)
1543
+
1544
+ # Same thing but boxing other
1545
+ other = tm.box_expected(other, box_with_array)
1546
+ if box_with_array is pd.array and op is roperator.radd:
1547
+ # We expect a NumpyExtensionArray, not ndarray[object] here
1548
+ expected = pd.array(expected, dtype=object)
1549
+ with tm.assert_produces_warning(performance_warning):
1550
+ res = op(dtarr, other)
1551
+ tm.assert_equal(res, expected)
1552
+
1553
+ @pytest.mark.parametrize(
1554
+ "op, offset, exp, exp_freq",
1555
+ [
1556
+ (
1557
+ "__add__",
1558
+ DateOffset(months=3, days=10),
1559
+ [
1560
+ Timestamp("2014-04-11"),
1561
+ Timestamp("2015-04-11"),
1562
+ Timestamp("2016-04-11"),
1563
+ Timestamp("2017-04-11"),
1564
+ ],
1565
+ None,
1566
+ ),
1567
+ (
1568
+ "__add__",
1569
+ DateOffset(months=3),
1570
+ [
1571
+ Timestamp("2014-04-01"),
1572
+ Timestamp("2015-04-01"),
1573
+ Timestamp("2016-04-01"),
1574
+ Timestamp("2017-04-01"),
1575
+ ],
1576
+ "YS-APR",
1577
+ ),
1578
+ (
1579
+ "__sub__",
1580
+ DateOffset(months=3, days=10),
1581
+ [
1582
+ Timestamp("2013-09-21"),
1583
+ Timestamp("2014-09-21"),
1584
+ Timestamp("2015-09-21"),
1585
+ Timestamp("2016-09-21"),
1586
+ ],
1587
+ None,
1588
+ ),
1589
+ (
1590
+ "__sub__",
1591
+ DateOffset(months=3),
1592
+ [
1593
+ Timestamp("2013-10-01"),
1594
+ Timestamp("2014-10-01"),
1595
+ Timestamp("2015-10-01"),
1596
+ Timestamp("2016-10-01"),
1597
+ ],
1598
+ "YS-OCT",
1599
+ ),
1600
+ ],
1601
+ )
1602
+ def test_dti_add_sub_nonzero_mth_offset(
1603
+ self, op, offset, exp, exp_freq, tz_aware_fixture, box_with_array
1604
+ ):
1605
+ # GH 26258
1606
+ tz = tz_aware_fixture
1607
+ date = date_range(
1608
+ start="01 Jan 2014", end="01 Jan 2017", freq="YS", tz=tz, unit="ns"
1609
+ )
1610
+ date = tm.box_expected(date, box_with_array, False)
1611
+ mth = getattr(date, op)
1612
+ result = mth(offset)
1613
+
1614
+ expected = DatetimeIndex(exp, tz=tz).as_unit("ns")
1615
+ expected = tm.box_expected(expected, box_with_array, False)
1616
+ tm.assert_equal(result, expected)
1617
+
1618
+ def test_dt64arr_series_add_DateOffset_with_milli(self):
1619
+ # GH 57529
1620
+ dti = DatetimeIndex(
1621
+ [
1622
+ "2000-01-01 00:00:00.012345678",
1623
+ "2000-01-31 00:00:00.012345678",
1624
+ "2000-02-29 00:00:00.012345678",
1625
+ ],
1626
+ dtype="datetime64[ns]",
1627
+ )
1628
+ result = dti + DateOffset(milliseconds=4)
1629
+ expected = DatetimeIndex(
1630
+ [
1631
+ "2000-01-01 00:00:00.016345678",
1632
+ "2000-01-31 00:00:00.016345678",
1633
+ "2000-02-29 00:00:00.016345678",
1634
+ ],
1635
+ dtype="datetime64[ns]",
1636
+ )
1637
+ tm.assert_index_equal(result, expected)
1638
+
1639
+ result = dti + DateOffset(days=1, milliseconds=4)
1640
+ expected = DatetimeIndex(
1641
+ [
1642
+ "2000-01-02 00:00:00.016345678",
1643
+ "2000-02-01 00:00:00.016345678",
1644
+ "2000-03-01 00:00:00.016345678",
1645
+ ],
1646
+ dtype="datetime64[ns]",
1647
+ )
1648
+ tm.assert_index_equal(result, expected)
1649
+
1650
+
1651
+ class TestDatetime64OverflowHandling:
1652
+ # TODO: box + de-duplicate
1653
+
1654
+ def test_dt64_overflow_masking(self, box_with_array):
1655
+ # GH#25317
1656
+ left = Series([Timestamp("1969-12-31")], dtype="M8[ns]")
1657
+ right = Series([NaT])
1658
+
1659
+ left = tm.box_expected(left, box_with_array)
1660
+ right = tm.box_expected(right, box_with_array)
1661
+
1662
+ expected = TimedeltaIndex([NaT], dtype="m8[ns]")
1663
+ expected = tm.box_expected(expected, box_with_array)
1664
+
1665
+ result = left - right
1666
+ tm.assert_equal(result, expected)
1667
+
1668
+ def test_dt64_series_arith_overflow(self):
1669
+ # GH#12534, fixed by GH#19024
1670
+ dt = Timestamp("1700-01-31")
1671
+ td = Timedelta("20000 Days")
1672
+ dti = date_range("1949-09-30", freq="100YE", periods=4, unit="ns")
1673
+ ser = Series(dti)
1674
+ msg = "Overflow in int64 addition"
1675
+ with pytest.raises(OverflowError, match=msg):
1676
+ ser - dt
1677
+ with pytest.raises(OverflowError, match=msg):
1678
+ dt - ser
1679
+ with pytest.raises(OverflowError, match=msg):
1680
+ ser + td
1681
+ with pytest.raises(OverflowError, match=msg):
1682
+ td + ser
1683
+
1684
+ ser.iloc[-1] = NaT
1685
+ expected = Series(
1686
+ ["2004-10-03", "2104-10-04", "2204-10-04", "NaT"], dtype="datetime64[ns]"
1687
+ )
1688
+ res = ser + td
1689
+ tm.assert_series_equal(res, expected)
1690
+ res = td + ser
1691
+ tm.assert_series_equal(res, expected)
1692
+
1693
+ ser.iloc[1:] = NaT
1694
+ expected = Series(["91279 Days", "NaT", "NaT", "NaT"], dtype="timedelta64[ns]")
1695
+ res = ser - dt
1696
+ tm.assert_series_equal(res, expected)
1697
+ res = dt - ser
1698
+ tm.assert_series_equal(res, -expected)
1699
+
1700
+ def test_datetimeindex_sub_timestamp_overflow(self):
1701
+ dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max]).as_unit("ns")
1702
+ dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min]).as_unit("ns")
1703
+
1704
+ tsneg = Timestamp("1950-01-01").as_unit("ns")
1705
+ ts_neg_variants = [
1706
+ tsneg,
1707
+ tsneg.to_pydatetime(),
1708
+ tsneg.to_datetime64().astype("datetime64[ns]"),
1709
+ tsneg.to_datetime64().astype("datetime64[D]"),
1710
+ ]
1711
+
1712
+ tspos = Timestamp("1980-01-01").as_unit("ns")
1713
+ ts_pos_variants = [
1714
+ tspos,
1715
+ tspos.to_pydatetime(),
1716
+ tspos.to_datetime64().astype("datetime64[ns]"),
1717
+ tspos.to_datetime64().astype("datetime64[D]"),
1718
+ ]
1719
+ msg = "Overflow in int64 addition"
1720
+ for variant in ts_neg_variants:
1721
+ with pytest.raises(OverflowError, match=msg):
1722
+ dtimax - variant
1723
+
1724
+ expected = Timestamp.max._value - tspos._value
1725
+ for variant in ts_pos_variants:
1726
+ res = dtimax - variant
1727
+ assert res[1]._value == expected
1728
+
1729
+ expected = Timestamp.min._value - tsneg._value
1730
+ for variant in ts_neg_variants:
1731
+ res = dtimin - variant
1732
+ assert res[1]._value == expected
1733
+
1734
+ for variant in ts_pos_variants:
1735
+ with pytest.raises(OverflowError, match=msg):
1736
+ dtimin - variant
1737
+
1738
+ def test_datetimeindex_sub_datetimeindex_overflow(self):
1739
+ # GH#22492, GH#22508
1740
+ dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max]).as_unit("ns")
1741
+ dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min]).as_unit("ns")
1742
+
1743
+ ts_neg = pd.to_datetime(["1950-01-01", "1950-01-01"]).as_unit("ns")
1744
+ ts_pos = pd.to_datetime(["1980-01-01", "1980-01-01"]).as_unit("ns")
1745
+
1746
+ # General tests
1747
+ expected = Timestamp.max._value - ts_pos[1]._value
1748
+ result = dtimax - ts_pos
1749
+ assert result[1]._value == expected
1750
+
1751
+ expected = Timestamp.min._value - ts_neg[1]._value
1752
+ result = dtimin - ts_neg
1753
+ assert result[1]._value == expected
1754
+ msg = "Overflow in int64 addition"
1755
+ with pytest.raises(OverflowError, match=msg):
1756
+ dtimax - ts_neg
1757
+
1758
+ with pytest.raises(OverflowError, match=msg):
1759
+ dtimin - ts_pos
1760
+
1761
+ # Edge cases
1762
+ tmin = pd.to_datetime([Timestamp.min])
1763
+ t1 = tmin + Timedelta.max + Timedelta("1us")
1764
+ with pytest.raises(OverflowError, match=msg):
1765
+ t1 - tmin
1766
+
1767
+ tmax = pd.to_datetime([Timestamp.max])
1768
+ t2 = tmax + Timedelta.min - Timedelta("1us")
1769
+ with pytest.raises(OverflowError, match=msg):
1770
+ tmax - t2
1771
+
1772
+
1773
+ class TestTimestampSeriesArithmetic:
1774
+ def test_empty_series_add_sub(self, box_with_array):
1775
+ # GH#13844
1776
+ a = Series(dtype="M8[ns]")
1777
+ b = Series(dtype="m8[ns]")
1778
+ a = box_with_array(a)
1779
+ b = box_with_array(b)
1780
+ tm.assert_equal(a, a + b)
1781
+ tm.assert_equal(a, a - b)
1782
+ tm.assert_equal(a, b + a)
1783
+ msg = "cannot subtract"
1784
+ with pytest.raises(TypeError, match=msg):
1785
+ b - a
1786
+
1787
+ def test_operators_datetimelike(self):
1788
+ # ## timedelta64 ###
1789
+ td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
1790
+ td1.iloc[2] = np.nan
1791
+
1792
+ # ## datetime64 ###
1793
+ dt1 = Series(
1794
+ [
1795
+ Timestamp("20111230"),
1796
+ Timestamp("20120101"),
1797
+ Timestamp("20120103"),
1798
+ ]
1799
+ )
1800
+ dt1.iloc[2] = np.nan
1801
+ dt2 = Series(
1802
+ [
1803
+ Timestamp("20111231"),
1804
+ Timestamp("20120102"),
1805
+ Timestamp("20120104"),
1806
+ ]
1807
+ )
1808
+ dt1 - dt2
1809
+ dt2 - dt1
1810
+
1811
+ # datetime64 with timetimedelta
1812
+ dt1 + td1
1813
+ td1 + dt1
1814
+ dt1 - td1
1815
+
1816
+ # timetimedelta with datetime64
1817
+ td1 + dt1
1818
+ dt1 + td1
1819
+
1820
+ def test_dt64ser_sub_datetime_dtype(self, unit):
1821
+ ts = Timestamp(datetime(1993, 1, 7, 13, 30, 00))
1822
+ dt = datetime(1993, 6, 22, 13, 30)
1823
+ ser = Series([ts], dtype=f"M8[{unit}]")
1824
+ result = ser - dt
1825
+
1826
+ # the expected unit is the max of `unit` and the unit imputed to `dt`,
1827
+ # which is "us"
1828
+ exp_unit = tm.get_finest_unit(unit, "us")
1829
+ assert result.dtype == f"timedelta64[{exp_unit}]"
1830
+
1831
+ # -------------------------------------------------------------
1832
+ # TODO: This next block of tests came from tests.series.test_operators,
1833
+ # needs to be de-duplicated and parametrized over `box` classes
1834
+
1835
+ @pytest.mark.parametrize(
1836
+ "left, right, op_fail",
1837
+ [
1838
+ [
1839
+ [Timestamp("20111230"), Timestamp("20120101"), NaT],
1840
+ [Timestamp("20111231"), Timestamp("20120102"), Timestamp("20120104")],
1841
+ ["__sub__", "__rsub__"],
1842
+ ],
1843
+ [
1844
+ [Timestamp("20111230"), Timestamp("20120101"), NaT],
1845
+ [timedelta(minutes=5, seconds=3), timedelta(minutes=5, seconds=3), NaT],
1846
+ ["__add__", "__radd__", "__sub__"],
1847
+ ],
1848
+ [
1849
+ [
1850
+ Timestamp("20111230", tz="US/Eastern"),
1851
+ Timestamp("20111230", tz="US/Eastern"),
1852
+ NaT,
1853
+ ],
1854
+ [timedelta(minutes=5, seconds=3), NaT, timedelta(minutes=5, seconds=3)],
1855
+ ["__add__", "__radd__", "__sub__"],
1856
+ ],
1857
+ ],
1858
+ )
1859
+ def test_operators_datetimelike_invalid(
1860
+ self, left, right, op_fail, all_arithmetic_operators
1861
+ ):
1862
+ # these are all TypeError ops
1863
+ op_str = all_arithmetic_operators
1864
+ arg1 = Series(left)
1865
+ arg2 = Series(right)
1866
+ # check that we are getting a TypeError
1867
+ # with 'operate' (from core/ops.py) for the ops that are not
1868
+ # defined
1869
+ op = getattr(arg1, op_str, None)
1870
+ # Previously, _validate_for_numeric_binop in core/indexes/base.py
1871
+ # did this for us.
1872
+ if op_str not in op_fail:
1873
+ with pytest.raises(
1874
+ TypeError, match="operate|[cC]annot|unsupported operand"
1875
+ ):
1876
+ op(arg2)
1877
+ else:
1878
+ # Smoke test
1879
+ op(arg2)
1880
+
1881
+ def test_sub_single_tz(self, unit):
1882
+ # GH#12290
1883
+ s1 = Series([Timestamp("2016-02-10", tz="America/Sao_Paulo")]).dt.as_unit(unit)
1884
+ s2 = Series([Timestamp("2016-02-08", tz="America/Sao_Paulo")]).dt.as_unit(unit)
1885
+ result = s1 - s2
1886
+ expected = Series([Timedelta("2days")]).dt.as_unit(unit)
1887
+ tm.assert_series_equal(result, expected)
1888
+ result = s2 - s1
1889
+ expected = Series([Timedelta("-2days")]).dt.as_unit(unit)
1890
+ tm.assert_series_equal(result, expected)
1891
+
1892
+ def test_dt64tz_series_sub_dtitz(self):
1893
+ # GH#19071 subtracting tzaware DatetimeIndex from tzaware Series
1894
+ # (with same tz) raises, fixed by #19024
1895
+ dti = date_range("1999-09-30", periods=10, tz="US/Pacific")
1896
+ ser = Series(dti)
1897
+ expected = Series(TimedeltaIndex(["0days"] * 10))
1898
+
1899
+ res = dti - ser
1900
+ tm.assert_series_equal(res, expected)
1901
+ res = ser - dti
1902
+ tm.assert_series_equal(res, expected)
1903
+
1904
+ def test_sub_datetime_compat(self, unit):
1905
+ # see GH#14088
1906
+ ser = Series([datetime(2016, 8, 23, 12, tzinfo=timezone.utc), NaT]).dt.as_unit(
1907
+ unit
1908
+ )
1909
+ dt = datetime(2016, 8, 22, 12, tzinfo=timezone.utc)
1910
+ # The datetime object has "us" so we upcast lower units
1911
+ exp_unit = tm.get_finest_unit(unit, "us")
1912
+ exp = Series([Timedelta("1 days"), NaT]).dt.as_unit(exp_unit)
1913
+ result = ser - dt
1914
+ tm.assert_series_equal(result, exp)
1915
+ result2 = ser - Timestamp(dt)
1916
+ tm.assert_series_equal(result2, exp)
1917
+
1918
+ def test_dt64_series_add_mixed_tick_DateOffset(self):
1919
+ # GH#4532
1920
+ # operate with pd.offsets
1921
+ s = Series([Timestamp("20130101 9:01"), Timestamp("20130101 9:02")])
1922
+
1923
+ result = s + pd.offsets.Milli(5)
1924
+ result2 = pd.offsets.Milli(5) + s
1925
+ expected = Series(
1926
+ [Timestamp("20130101 9:01:00.005"), Timestamp("20130101 9:02:00.005")]
1927
+ )
1928
+ tm.assert_series_equal(result, expected)
1929
+ tm.assert_series_equal(result2, expected)
1930
+
1931
+ result = s + pd.offsets.Minute(5) + pd.offsets.Milli(5)
1932
+ expected = Series(
1933
+ [Timestamp("20130101 9:06:00.005"), Timestamp("20130101 9:07:00.005")]
1934
+ )
1935
+ tm.assert_series_equal(result, expected)
1936
+
1937
+ def test_datetime64_ops_nat(self, unit):
1938
+ # GH#11349
1939
+ datetime_series = Series([NaT, Timestamp("19900315")]).dt.as_unit(unit)
1940
+ nat_series_dtype_timestamp = Series([NaT, NaT], dtype=f"datetime64[{unit}]")
1941
+ single_nat_dtype_datetime = Series([NaT], dtype=f"datetime64[{unit}]")
1942
+
1943
+ # subtraction
1944
+ tm.assert_series_equal(-NaT + datetime_series, nat_series_dtype_timestamp)
1945
+ msg = "bad operand type for unary -: 'DatetimeArray'"
1946
+ with pytest.raises(TypeError, match=msg):
1947
+ -single_nat_dtype_datetime + datetime_series
1948
+
1949
+ tm.assert_series_equal(
1950
+ -NaT + nat_series_dtype_timestamp, nat_series_dtype_timestamp
1951
+ )
1952
+ with pytest.raises(TypeError, match=msg):
1953
+ -single_nat_dtype_datetime + nat_series_dtype_timestamp
1954
+
1955
+ # addition
1956
+ tm.assert_series_equal(
1957
+ nat_series_dtype_timestamp + NaT, nat_series_dtype_timestamp
1958
+ )
1959
+ tm.assert_series_equal(
1960
+ NaT + nat_series_dtype_timestamp, nat_series_dtype_timestamp
1961
+ )
1962
+
1963
+ tm.assert_series_equal(
1964
+ nat_series_dtype_timestamp + NaT, nat_series_dtype_timestamp
1965
+ )
1966
+ tm.assert_series_equal(
1967
+ NaT + nat_series_dtype_timestamp, nat_series_dtype_timestamp
1968
+ )
1969
+
1970
+ # -------------------------------------------------------------
1971
+ # Timezone-Centric Tests
1972
+
1973
+ def test_operators_datetimelike_with_timezones(self):
1974
+ tz = "US/Eastern"
1975
+ dt1 = Series(date_range("2000-01-01 09:00:00", periods=5, tz=tz), name="foo")
1976
+ dt2 = dt1.copy()
1977
+ dt2.iloc[2] = np.nan
1978
+
1979
+ td1 = Series(pd.timedelta_range("1 days 1 min", periods=5, freq="h"))
1980
+ td2 = td1.copy()
1981
+ td2.iloc[1] = np.nan
1982
+ assert td2._values.freq is None
1983
+
1984
+ result = dt1 + td1[0]
1985
+ exp = (dt1.dt.tz_localize(None) + td1[0]).dt.tz_localize(tz)
1986
+ tm.assert_series_equal(result, exp)
1987
+
1988
+ result = dt2 + td2[0]
1989
+ exp = (dt2.dt.tz_localize(None) + td2[0]).dt.tz_localize(tz)
1990
+ tm.assert_series_equal(result, exp)
1991
+
1992
+ # odd numpy behavior with scalar timedeltas
1993
+ result = td1[0] + dt1
1994
+ exp = (dt1.dt.tz_localize(None) + td1[0]).dt.tz_localize(tz)
1995
+ tm.assert_series_equal(result, exp)
1996
+
1997
+ result = td2[0] + dt2
1998
+ exp = (dt2.dt.tz_localize(None) + td2[0]).dt.tz_localize(tz)
1999
+ tm.assert_series_equal(result, exp)
2000
+
2001
+ result = dt1 - td1[0]
2002
+ exp = (dt1.dt.tz_localize(None) - td1[0]).dt.tz_localize(tz)
2003
+ tm.assert_series_equal(result, exp)
2004
+ msg = "cannot subtract DatetimeArray from"
2005
+ with pytest.raises(TypeError, match=msg):
2006
+ td1[0] - dt1
2007
+
2008
+ result = dt2 - td2[0]
2009
+ exp = (dt2.dt.tz_localize(None) - td2[0]).dt.tz_localize(tz)
2010
+ tm.assert_series_equal(result, exp)
2011
+ with pytest.raises(TypeError, match=msg):
2012
+ td2[0] - dt2
2013
+
2014
+ result = dt1 + td1
2015
+ exp = (dt1.dt.tz_localize(None) + td1).dt.tz_localize(tz)
2016
+ tm.assert_series_equal(result, exp)
2017
+
2018
+ result = dt2 + td2
2019
+ exp = (dt2.dt.tz_localize(None) + td2).dt.tz_localize(tz)
2020
+ tm.assert_series_equal(result, exp)
2021
+
2022
+ result = dt1 - td1
2023
+ exp = (dt1.dt.tz_localize(None) - td1).dt.tz_localize(tz)
2024
+ tm.assert_series_equal(result, exp)
2025
+
2026
+ result = dt2 - td2
2027
+ exp = (dt2.dt.tz_localize(None) - td2).dt.tz_localize(tz)
2028
+ tm.assert_series_equal(result, exp)
2029
+ msg = "cannot (add|subtract)"
2030
+ with pytest.raises(TypeError, match=msg):
2031
+ td1 - dt1
2032
+ with pytest.raises(TypeError, match=msg):
2033
+ td2 - dt2
2034
+
2035
+
2036
+ class TestDatetimeIndexArithmetic:
2037
+ # -------------------------------------------------------------
2038
+ # Binary operations DatetimeIndex and TimedeltaIndex/array
2039
+
2040
+ def test_dti_add_tdi(self, tz_naive_fixture):
2041
+ # GH#17558
2042
+ tz = tz_naive_fixture
2043
+ dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
2044
+ tdi = pd.timedelta_range("0 days", periods=10)
2045
+ expected = date_range("2017-01-01", periods=10, tz=tz)
2046
+ expected = expected._with_freq(None)
2047
+
2048
+ # add with TimedeltaIndex
2049
+ result = dti + tdi
2050
+ tm.assert_index_equal(result, expected)
2051
+
2052
+ result = tdi + dti
2053
+ tm.assert_index_equal(result, expected)
2054
+
2055
+ # add with timedelta64 array
2056
+ result = dti + tdi.values
2057
+ tm.assert_index_equal(result, expected)
2058
+
2059
+ result = tdi.values + dti
2060
+ tm.assert_index_equal(result, expected)
2061
+
2062
+ def test_dti_iadd_tdi(self, tz_naive_fixture):
2063
+ # GH#17558
2064
+ tz = tz_naive_fixture
2065
+ dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
2066
+ tdi = pd.timedelta_range("0 days", periods=10)
2067
+ expected = date_range("2017-01-01", periods=10, tz=tz)
2068
+ expected = expected._with_freq(None)
2069
+
2070
+ # iadd with TimedeltaIndex
2071
+ result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
2072
+ result += tdi
2073
+ tm.assert_index_equal(result, expected)
2074
+
2075
+ result = pd.timedelta_range("0 days", periods=10)
2076
+ result += dti
2077
+ tm.assert_index_equal(result, expected)
2078
+
2079
+ # iadd with timedelta64 array
2080
+ result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
2081
+ result += tdi.values
2082
+ tm.assert_index_equal(result, expected)
2083
+
2084
+ result = pd.timedelta_range("0 days", periods=10)
2085
+ result += dti
2086
+ tm.assert_index_equal(result, expected)
2087
+
2088
+ def test_dti_sub_tdi(self, tz_naive_fixture):
2089
+ # GH#17558
2090
+ tz = tz_naive_fixture
2091
+ dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
2092
+ tdi = pd.timedelta_range("0 days", periods=10)
2093
+ expected = date_range("2017-01-01", periods=10, tz=tz, freq="-1D")
2094
+ expected = expected._with_freq(None)
2095
+
2096
+ # sub with TimedeltaIndex
2097
+ result = dti - tdi
2098
+ tm.assert_index_equal(result, expected)
2099
+
2100
+ msg = "cannot subtract .*TimedeltaArray"
2101
+ with pytest.raises(TypeError, match=msg):
2102
+ tdi - dti
2103
+
2104
+ # sub with timedelta64 array
2105
+ result = dti - tdi.values
2106
+ tm.assert_index_equal(result, expected)
2107
+
2108
+ msg = "cannot subtract a datelike from a TimedeltaArray"
2109
+ with pytest.raises(TypeError, match=msg):
2110
+ tdi.values - dti
2111
+
2112
+ def test_dti_isub_tdi(self, tz_naive_fixture, unit):
2113
+ # GH#17558
2114
+ tz = tz_naive_fixture
2115
+ dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
2116
+ tdi = pd.timedelta_range("0 days", periods=10, unit=unit)
2117
+ expected = date_range("2017-01-01", periods=10, tz=tz, freq="-1D", unit=unit)
2118
+ expected = expected._with_freq(None)
2119
+
2120
+ # isub with TimedeltaIndex
2121
+ result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
2122
+ result -= tdi
2123
+ tm.assert_index_equal(result, expected)
2124
+
2125
+ # DTA.__isub__ GH#43904
2126
+ dta = dti._data.copy()
2127
+ dta -= tdi
2128
+ tm.assert_datetime_array_equal(dta, expected._data)
2129
+
2130
+ out = dti._data.copy()
2131
+ np.subtract(out, tdi, out=out)
2132
+ tm.assert_datetime_array_equal(out, expected._data)
2133
+
2134
+ msg = "cannot subtract a datelike from a TimedeltaArray"
2135
+ with pytest.raises(TypeError, match=msg):
2136
+ tdi -= dti
2137
+
2138
+ # isub with timedelta64 array
2139
+ result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
2140
+ result -= tdi.values
2141
+ tm.assert_index_equal(result, expected)
2142
+
2143
+ with pytest.raises(TypeError, match=msg):
2144
+ tdi.values -= dti
2145
+
2146
+ with pytest.raises(TypeError, match=msg):
2147
+ tdi._values -= dti
2148
+
2149
+ # -------------------------------------------------------------
2150
+ # Binary Operations DatetimeIndex and datetime-like
2151
+ # TODO: A couple other tests belong in this section. Move them in
2152
+ # A PR where there isn't already a giant diff.
2153
+
2154
+ # -------------------------------------------------------------
2155
+
2156
+ def test_dta_add_sub_index(self, tz_naive_fixture):
2157
+ # Check that DatetimeArray defers to Index classes
2158
+ dti = date_range("20130101", periods=3, tz=tz_naive_fixture)
2159
+ dta = dti.array
2160
+ result = dta - dti
2161
+ expected = dti - dti
2162
+ tm.assert_index_equal(result, expected)
2163
+
2164
+ tdi = result
2165
+ result = dta + tdi
2166
+ expected = dti + tdi
2167
+ tm.assert_index_equal(result, expected)
2168
+
2169
+ result = dta - tdi
2170
+ expected = dti - tdi
2171
+ tm.assert_index_equal(result, expected)
2172
+
2173
+ def test_sub_dti_dti(self, unit):
2174
+ # previously performed setop (deprecated in 0.16.0), now changed to
2175
+ # return subtraction -> TimeDeltaIndex (GH ...)
2176
+
2177
+ dti = date_range("20130101", periods=3, unit=unit)
2178
+ dti_tz = date_range("20130101", periods=3, unit=unit).tz_localize("US/Eastern")
2179
+ expected = TimedeltaIndex([0, 0, 0]).as_unit(unit)
2180
+
2181
+ result = dti - dti
2182
+ tm.assert_index_equal(result, expected)
2183
+
2184
+ result = dti_tz - dti_tz
2185
+ tm.assert_index_equal(result, expected)
2186
+ msg = "Cannot subtract tz-naive and tz-aware datetime-like objects"
2187
+ with pytest.raises(TypeError, match=msg):
2188
+ dti_tz - dti
2189
+
2190
+ with pytest.raises(TypeError, match=msg):
2191
+ dti - dti_tz
2192
+
2193
+ # isub
2194
+ dti -= dti
2195
+ tm.assert_index_equal(dti, expected)
2196
+
2197
+ # different length raises ValueError
2198
+ dti1 = date_range("20130101", periods=3, unit=unit)
2199
+ dti2 = date_range("20130101", periods=4, unit=unit)
2200
+ msg = "cannot add indices of unequal length"
2201
+ with pytest.raises(ValueError, match=msg):
2202
+ dti1 - dti2
2203
+
2204
+ # NaN propagation
2205
+ dti1 = DatetimeIndex(["2012-01-01", np.nan, "2012-01-03"]).as_unit(unit)
2206
+ dti2 = DatetimeIndex(["2012-01-02", "2012-01-03", np.nan]).as_unit(unit)
2207
+ expected = TimedeltaIndex(["1 days", np.nan, np.nan]).as_unit(unit)
2208
+ result = dti2 - dti1
2209
+ tm.assert_index_equal(result, expected)
2210
+
2211
+ # -------------------------------------------------------------------
2212
+ # TODO: Most of this block is moved from series or frame tests, needs
2213
+ # cleanup, box-parametrization, and de-duplication
2214
+
2215
+ @pytest.mark.parametrize("op", [operator.add, operator.sub])
2216
+ def test_timedelta64_equal_timedelta_supported_ops(self, op, box_with_array):
2217
+ ser = Series(
2218
+ [
2219
+ Timestamp("20130301"),
2220
+ Timestamp("20130228 23:00:00"),
2221
+ Timestamp("20130228 22:00:00"),
2222
+ Timestamp("20130228 21:00:00"),
2223
+ ]
2224
+ )
2225
+ obj = box_with_array(ser)
2226
+
2227
+ intervals = ["D", "h", "m", "s", "us"]
2228
+
2229
+ def timedelta64(*args):
2230
+ # see casting notes in NumPy gh-12927
2231
+ return np.sum(list(map(np.timedelta64, args, intervals)))
2232
+
2233
+ for d, h, m, s, us in product(*([range(2)] * 5)):
2234
+ nptd = timedelta64(d, h, m, s, us)
2235
+ pytd = timedelta(days=d, hours=h, minutes=m, seconds=s, microseconds=us)
2236
+ lhs = op(obj, nptd)
2237
+ rhs = op(obj, pytd)
2238
+
2239
+ tm.assert_equal(lhs, rhs)
2240
+
2241
+ def test_ops_nat_mixed_datetime64_timedelta64(self):
2242
+ # GH#11349
2243
+ timedelta_series = Series([NaT, Timedelta("1s")])
2244
+ datetime_series = Series([NaT, Timestamp("19900315")])
2245
+ nat_series_dtype_timedelta = Series([NaT, NaT], dtype="timedelta64[ns]")
2246
+ nat_series_dtype_timestamp = Series([NaT, NaT], dtype="datetime64[ns]")
2247
+ single_nat_dtype_datetime = Series([NaT], dtype="datetime64[ns]")
2248
+ single_nat_dtype_timedelta = Series([NaT], dtype="timedelta64[ns]")
2249
+
2250
+ # subtraction
2251
+ tm.assert_series_equal(
2252
+ datetime_series - single_nat_dtype_datetime, nat_series_dtype_timedelta
2253
+ )
2254
+
2255
+ tm.assert_series_equal(
2256
+ datetime_series - single_nat_dtype_timedelta, nat_series_dtype_timestamp
2257
+ )
2258
+ tm.assert_series_equal(
2259
+ -single_nat_dtype_timedelta + datetime_series, nat_series_dtype_timestamp
2260
+ )
2261
+
2262
+ # without a Series wrapping the NaT, it is ambiguous
2263
+ # whether it is a datetime64 or timedelta64
2264
+ # defaults to interpreting it as timedelta64
2265
+ tm.assert_series_equal(
2266
+ nat_series_dtype_timestamp - single_nat_dtype_datetime,
2267
+ nat_series_dtype_timedelta,
2268
+ )
2269
+
2270
+ tm.assert_series_equal(
2271
+ nat_series_dtype_timestamp - single_nat_dtype_timedelta,
2272
+ nat_series_dtype_timestamp,
2273
+ )
2274
+ tm.assert_series_equal(
2275
+ -single_nat_dtype_timedelta + nat_series_dtype_timestamp,
2276
+ nat_series_dtype_timestamp,
2277
+ )
2278
+ msg = "cannot subtract a datelike"
2279
+ with pytest.raises(TypeError, match=msg):
2280
+ timedelta_series - single_nat_dtype_datetime
2281
+
2282
+ # addition
2283
+ tm.assert_series_equal(
2284
+ nat_series_dtype_timestamp + single_nat_dtype_timedelta,
2285
+ nat_series_dtype_timestamp,
2286
+ )
2287
+ tm.assert_series_equal(
2288
+ single_nat_dtype_timedelta + nat_series_dtype_timestamp,
2289
+ nat_series_dtype_timestamp,
2290
+ )
2291
+
2292
+ tm.assert_series_equal(
2293
+ nat_series_dtype_timestamp + single_nat_dtype_timedelta,
2294
+ nat_series_dtype_timestamp,
2295
+ )
2296
+ tm.assert_series_equal(
2297
+ single_nat_dtype_timedelta + nat_series_dtype_timestamp,
2298
+ nat_series_dtype_timestamp,
2299
+ )
2300
+
2301
+ tm.assert_series_equal(
2302
+ nat_series_dtype_timedelta + single_nat_dtype_datetime,
2303
+ nat_series_dtype_timestamp,
2304
+ )
2305
+ tm.assert_series_equal(
2306
+ single_nat_dtype_datetime + nat_series_dtype_timedelta,
2307
+ nat_series_dtype_timestamp,
2308
+ )
2309
+
2310
+ def test_ufunc_coercions(self, unit):
2311
+ idx = date_range("2011-01-01", periods=3, freq="2D", name="x", unit=unit)
2312
+
2313
+ delta = np.timedelta64(1, "D")
2314
+ exp = date_range("2011-01-02", periods=3, freq="2D", name="x", unit=unit)
2315
+ for result in [idx + delta, np.add(idx, delta)]:
2316
+ assert isinstance(result, DatetimeIndex)
2317
+ tm.assert_index_equal(result, exp)
2318
+ assert result.freq == "2D"
2319
+
2320
+ exp = date_range("2010-12-31", periods=3, freq="2D", name="x", unit=unit)
2321
+
2322
+ for result in [idx - delta, np.subtract(idx, delta)]:
2323
+ assert isinstance(result, DatetimeIndex)
2324
+ tm.assert_index_equal(result, exp)
2325
+ assert result.freq == "2D"
2326
+
2327
+ # When adding/subtracting an ndarray (which has no .freq), the result
2328
+ # does not infer freq
2329
+ idx = idx._with_freq(None)
2330
+ delta = np.array(
2331
+ [np.timedelta64(1, "D"), np.timedelta64(2, "D"), np.timedelta64(3, "D")]
2332
+ )
2333
+ exp = DatetimeIndex(
2334
+ ["2011-01-02", "2011-01-05", "2011-01-08"], name="x"
2335
+ ).as_unit(unit)
2336
+
2337
+ for result in [idx + delta, np.add(idx, delta)]:
2338
+ tm.assert_index_equal(result, exp)
2339
+ assert result.freq == exp.freq
2340
+
2341
+ exp = DatetimeIndex(
2342
+ ["2010-12-31", "2011-01-01", "2011-01-02"], name="x"
2343
+ ).as_unit(unit)
2344
+ for result in [idx - delta, np.subtract(idx, delta)]:
2345
+ assert isinstance(result, DatetimeIndex)
2346
+ tm.assert_index_equal(result, exp)
2347
+ assert result.freq == exp.freq
2348
+
2349
+ def test_dti_add_series(self, tz_naive_fixture, names):
2350
+ # GH#13905
2351
+ tz = tz_naive_fixture
2352
+ index = DatetimeIndex(
2353
+ ["2016-06-28 05:30", "2016-06-28 05:31"], tz=tz, name=names[0]
2354
+ ).as_unit("ns")
2355
+ ser = Series([Timedelta(seconds=5)] * 2, index=index, name=names[1])
2356
+ expected = Series(index + Timedelta(seconds=5), index=index, name=names[2])
2357
+
2358
+ # passing name arg isn't enough when names[2] is None
2359
+ expected.name = names[2]
2360
+ assert expected.dtype == index.dtype
2361
+ result = ser + index
2362
+ tm.assert_series_equal(result, expected)
2363
+ result2 = index + ser
2364
+ tm.assert_series_equal(result2, expected)
2365
+
2366
+ expected = index + Timedelta(seconds=5)
2367
+ result3 = ser.values + index
2368
+ tm.assert_index_equal(result3, expected)
2369
+ result4 = index + ser.values
2370
+ tm.assert_index_equal(result4, expected)
2371
+
2372
+ @pytest.mark.parametrize("op", [operator.add, roperator.radd, operator.sub])
2373
+ def test_dti_addsub_offset_arraylike(
2374
+ self, performance_warning, tz_naive_fixture, names, op, index_or_series
2375
+ ):
2376
+ # GH#18849, GH#19744
2377
+ other_box = index_or_series
2378
+
2379
+ tz = tz_naive_fixture
2380
+ dti = date_range("2017-01-01", periods=2, tz=tz, name=names[0])
2381
+ other = other_box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)], name=names[1])
2382
+
2383
+ xbox = get_upcast_box(dti, other)
2384
+
2385
+ with tm.assert_produces_warning(performance_warning):
2386
+ res = op(dti, other)
2387
+
2388
+ expected = DatetimeIndex(
2389
+ [op(dti[n], other[n]) for n in range(len(dti))], name=names[2], freq="infer"
2390
+ )
2391
+ expected = tm.box_expected(expected, xbox).astype(object)
2392
+ tm.assert_equal(res, expected)
2393
+
2394
+ @pytest.mark.parametrize("other_box", [pd.Index, np.array])
2395
+ def test_dti_addsub_object_arraylike(
2396
+ self, performance_warning, tz_naive_fixture, box_with_array, other_box
2397
+ ):
2398
+ tz = tz_naive_fixture
2399
+
2400
+ dti = date_range("2017-01-01", periods=2, tz=tz)
2401
+ dtarr = tm.box_expected(dti, box_with_array)
2402
+ other = other_box([pd.offsets.MonthEnd(), Timedelta(days=4)])
2403
+ xbox = get_upcast_box(dtarr, other)
2404
+
2405
+ expected = DatetimeIndex(["2017-01-31", "2017-01-06"], tz=tz_naive_fixture)
2406
+ expected = tm.box_expected(expected, xbox).astype(object)
2407
+
2408
+ with tm.assert_produces_warning(performance_warning):
2409
+ result = dtarr + other
2410
+ tm.assert_equal(result, expected)
2411
+
2412
+ expected = DatetimeIndex(["2016-12-31", "2016-12-29"], tz=tz_naive_fixture)
2413
+ expected = tm.box_expected(expected, xbox).astype(object)
2414
+
2415
+ with tm.assert_produces_warning(performance_warning):
2416
+ result = dtarr - other
2417
+ tm.assert_equal(result, expected)
2418
+
2419
+
2420
+ @pytest.mark.parametrize("years", [-1, 0, 1])
2421
+ @pytest.mark.parametrize("months", [-2, 0, 2])
2422
+ def test_shift_months(years, months, unit):
2423
+ dti = DatetimeIndex(
2424
+ [
2425
+ Timestamp("2000-01-05 00:15:00"),
2426
+ Timestamp("2000-01-31 00:23:00"),
2427
+ Timestamp("2000-01-01"),
2428
+ Timestamp("2000-02-29"),
2429
+ Timestamp("2000-12-31"),
2430
+ ]
2431
+ ).as_unit(unit)
2432
+ shifted = shift_months(dti.asi8, years * 12 + months, reso=dti._data._creso)
2433
+ shifted_dt64 = shifted.view(f"M8[{dti.unit}]")
2434
+ actual = DatetimeIndex(shifted_dt64)
2435
+
2436
+ raw = [x + pd.offsets.DateOffset(years=years, months=months) for x in dti]
2437
+ expected = DatetimeIndex(raw).as_unit(dti.unit)
2438
+ tm.assert_index_equal(actual, expected)
2439
+
2440
+
2441
+ def test_dt64arr_addsub_object_dtype_2d(performance_warning):
2442
+ # block-wise DataFrame operations will require operating on 2D
2443
+ # DatetimeArray/TimedeltaArray, so check that specifically.
2444
+ dti = date_range("1994-02-13", freq="2W", periods=4)
2445
+ dta = dti._data.reshape((4, 1))
2446
+
2447
+ other = np.array([[pd.offsets.Day(n)] for n in range(4)])
2448
+ assert other.shape == dta.shape
2449
+
2450
+ with tm.assert_produces_warning(performance_warning):
2451
+ result = dta + other
2452
+ with tm.assert_produces_warning(performance_warning):
2453
+ expected = (dta[:, 0] + other[:, 0]).reshape(-1, 1)
2454
+
2455
+ tm.assert_numpy_array_equal(result, expected)
2456
+
2457
+ with tm.assert_produces_warning(performance_warning):
2458
+ # Case where we expect to get a TimedeltaArray back
2459
+ result2 = dta - dta.astype(object)
2460
+
2461
+ assert result2.shape == (4, 1)
2462
+ assert all(td._value == 0 for td in result2.ravel())
2463
+
2464
+
2465
+ def test_non_nano_dt64_addsub_np_nat_scalars():
2466
+ # GH 52295
2467
+ ser = Series([1233242342344, 232432434324, 332434242344], dtype="datetime64[ms]")
2468
+ result = ser - np.datetime64("nat", "ms")
2469
+ expected = Series([NaT] * 3, dtype="timedelta64[ms]")
2470
+ tm.assert_series_equal(result, expected)
2471
+
2472
+ result = ser + np.timedelta64("nat", "ms")
2473
+ expected = Series([NaT] * 3, dtype="datetime64[ms]")
2474
+ tm.assert_series_equal(result, expected)
2475
+
2476
+
2477
+ def test_non_nano_dt64_addsub_np_nat_scalars_unitless():
2478
+ # GH 52295
2479
+ # TODO: Can we default to the ser unit?
2480
+ ser = Series([1233242342344, 232432434324, 332434242344], dtype="datetime64[ms]")
2481
+ result = ser - np.datetime64("nat", "ms")
2482
+ expected = Series([NaT] * 3, dtype="timedelta64[ms]")
2483
+ tm.assert_series_equal(result, expected)
2484
+
2485
+ result = ser + np.timedelta64("NaT", "ms")
2486
+ expected = Series([NaT] * 3, dtype="datetime64[ms]")
2487
+ tm.assert_series_equal(result, expected)
2488
+
2489
+
2490
+ def test_non_nano_dt64_addsub_np_nat_scalars_unsupported_unit():
2491
+ # GH 52295
2492
+ ser = Series([12332, 23243, 33243], dtype="datetime64[s]")
2493
+ result = ser - np.datetime64("nat", "D")
2494
+ expected = Series([NaT] * 3, dtype="timedelta64[s]")
2495
+ tm.assert_series_equal(result, expected)
2496
+
2497
+ result = ser + np.timedelta64("nat", "D")
2498
+ expected = Series([NaT] * 3, dtype="datetime64[s]")
2499
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_interval.py ADDED
@@ -0,0 +1,308 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import operator
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ from pandas.core.dtypes.common import is_list_like
7
+
8
+ import pandas as pd
9
+ from pandas import (
10
+ Categorical,
11
+ Index,
12
+ Interval,
13
+ IntervalIndex,
14
+ Period,
15
+ Series,
16
+ Timedelta,
17
+ Timestamp,
18
+ date_range,
19
+ period_range,
20
+ timedelta_range,
21
+ )
22
+ import pandas._testing as tm
23
+ from pandas.core.arrays import (
24
+ BooleanArray,
25
+ IntervalArray,
26
+ )
27
+ from pandas.tests.arithmetic.common import get_upcast_box
28
+
29
+
30
+ @pytest.fixture(
31
+ params=[
32
+ (Index([0, 2, 4, 4]), Index([1, 3, 5, 8])),
33
+ (Index([0.0, 1.0, 2.0, np.nan]), Index([1.0, 2.0, 3.0, np.nan])),
34
+ (
35
+ timedelta_range("0 days", periods=3).insert(3, pd.NaT),
36
+ timedelta_range("1 day", periods=3).insert(3, pd.NaT),
37
+ ),
38
+ (
39
+ date_range("20170101", periods=3).insert(3, pd.NaT),
40
+ date_range("20170102", periods=3).insert(3, pd.NaT),
41
+ ),
42
+ (
43
+ date_range("20170101", periods=3, tz="US/Eastern").insert(3, pd.NaT),
44
+ date_range("20170102", periods=3, tz="US/Eastern").insert(3, pd.NaT),
45
+ ),
46
+ ],
47
+ ids=lambda x: str(x[0].dtype),
48
+ )
49
+ def left_right_dtypes(request):
50
+ """
51
+ Fixture for building an IntervalArray from various dtypes
52
+ """
53
+ return request.param
54
+
55
+
56
+ @pytest.fixture
57
+ def interval_array(left_right_dtypes):
58
+ """
59
+ Fixture to generate an IntervalArray of various dtypes containing NA if possible
60
+ """
61
+ left, right = left_right_dtypes
62
+ return IntervalArray.from_arrays(left, right)
63
+
64
+
65
+ def create_categorical_intervals(left, right, closed="right"):
66
+ return Categorical(IntervalIndex.from_arrays(left, right, closed))
67
+
68
+
69
+ def create_series_intervals(left, right, closed="right"):
70
+ return Series(IntervalArray.from_arrays(left, right, closed))
71
+
72
+
73
+ def create_series_categorical_intervals(left, right, closed="right"):
74
+ return Series(Categorical(IntervalIndex.from_arrays(left, right, closed)))
75
+
76
+
77
+ class TestComparison:
78
+ @pytest.fixture(params=[operator.eq, operator.ne])
79
+ def op(self, request):
80
+ return request.param
81
+
82
+ @pytest.fixture(
83
+ params=[
84
+ IntervalArray.from_arrays,
85
+ IntervalIndex.from_arrays,
86
+ create_categorical_intervals,
87
+ create_series_intervals,
88
+ create_series_categorical_intervals,
89
+ ],
90
+ ids=[
91
+ "IntervalArray",
92
+ "IntervalIndex",
93
+ "Categorical[Interval]",
94
+ "Series[Interval]",
95
+ "Series[Categorical[Interval]]",
96
+ ],
97
+ )
98
+ def interval_constructor(self, request):
99
+ """
100
+ Fixture for all pandas native interval constructors.
101
+ To be used as the LHS of IntervalArray comparisons.
102
+ """
103
+ return request.param
104
+
105
+ def elementwise_comparison(self, op, interval_array, other):
106
+ """
107
+ Helper that performs elementwise comparisons between `array` and `other`
108
+ """
109
+ other = other if is_list_like(other) else [other] * len(interval_array)
110
+ expected = np.array(
111
+ [op(x, y) for x, y in zip(interval_array, other, strict=True)]
112
+ )
113
+ if isinstance(other, Series):
114
+ return Series(expected, index=other.index)
115
+ return expected
116
+
117
+ def test_compare_scalar_interval(self, op, interval_array):
118
+ # matches first interval
119
+ other = interval_array[0]
120
+ result = op(interval_array, other)
121
+ expected = self.elementwise_comparison(op, interval_array, other)
122
+ tm.assert_numpy_array_equal(result, expected)
123
+
124
+ # matches on a single endpoint but not both
125
+ other = Interval(interval_array.left[0], interval_array.right[1])
126
+ result = op(interval_array, other)
127
+ expected = self.elementwise_comparison(op, interval_array, other)
128
+ tm.assert_numpy_array_equal(result, expected)
129
+
130
+ def test_compare_scalar_interval_mixed_closed(self, op, closed, other_closed):
131
+ interval_array = IntervalArray.from_arrays(range(2), range(1, 3), closed=closed)
132
+ other = Interval(0, 1, closed=other_closed)
133
+
134
+ result = op(interval_array, other)
135
+ expected = self.elementwise_comparison(op, interval_array, other)
136
+ tm.assert_numpy_array_equal(result, expected)
137
+
138
+ def test_compare_scalar_na(self, op, interval_array, nulls_fixture, box_with_array):
139
+ box = box_with_array
140
+ obj = tm.box_expected(interval_array, box)
141
+ result = op(obj, nulls_fixture)
142
+
143
+ if nulls_fixture is pd.NA:
144
+ # GH#31882
145
+ exp = np.ones(interval_array.shape, dtype=bool)
146
+ expected = BooleanArray(exp, exp)
147
+ else:
148
+ expected = self.elementwise_comparison(op, interval_array, nulls_fixture)
149
+
150
+ if not (box is Index and nulls_fixture is pd.NA):
151
+ # don't cast expected from BooleanArray to ndarray[object]
152
+ xbox = get_upcast_box(obj, nulls_fixture, True)
153
+ expected = tm.box_expected(expected, xbox)
154
+
155
+ tm.assert_equal(result, expected)
156
+
157
+ rev = op(nulls_fixture, obj)
158
+ tm.assert_equal(rev, expected)
159
+
160
+ @pytest.mark.parametrize(
161
+ "other",
162
+ [
163
+ 0,
164
+ 1.0,
165
+ True,
166
+ "foo",
167
+ Timestamp("2017-01-01"),
168
+ Timestamp("2017-01-01", tz="US/Eastern"),
169
+ Timedelta("0 days"),
170
+ Period("2017-01-01", "D"),
171
+ ],
172
+ )
173
+ def test_compare_scalar_other(self, op, interval_array, other):
174
+ result = op(interval_array, other)
175
+ expected = self.elementwise_comparison(op, interval_array, other)
176
+ tm.assert_numpy_array_equal(result, expected)
177
+
178
+ def test_compare_list_like_interval(self, op, interval_array, interval_constructor):
179
+ # same endpoints
180
+ other = interval_constructor(interval_array.left, interval_array.right)
181
+ result = op(interval_array, other)
182
+ expected = self.elementwise_comparison(op, interval_array, other)
183
+ tm.assert_equal(result, expected)
184
+
185
+ # different endpoints
186
+ other = interval_constructor(
187
+ interval_array.left[::-1], interval_array.right[::-1]
188
+ )
189
+ result = op(interval_array, other)
190
+ expected = self.elementwise_comparison(op, interval_array, other)
191
+ tm.assert_equal(result, expected)
192
+
193
+ # all nan endpoints
194
+ other = interval_constructor([np.nan] * 4, [np.nan] * 4)
195
+ result = op(interval_array, other)
196
+ expected = self.elementwise_comparison(op, interval_array, other)
197
+ tm.assert_equal(result, expected)
198
+
199
+ def test_compare_list_like_interval_mixed_closed(
200
+ self, op, interval_constructor, closed, other_closed
201
+ ):
202
+ interval_array = IntervalArray.from_arrays(range(2), range(1, 3), closed=closed)
203
+ other = interval_constructor(range(2), range(1, 3), closed=other_closed)
204
+
205
+ result = op(interval_array, other)
206
+ expected = self.elementwise_comparison(op, interval_array, other)
207
+ tm.assert_equal(result, expected)
208
+
209
+ @pytest.mark.parametrize(
210
+ "other",
211
+ [
212
+ (
213
+ Interval(0, 1),
214
+ Interval(Timedelta("1 day"), Timedelta("2 days")),
215
+ Interval(4, 5, "both"),
216
+ Interval(10, 20, "neither"),
217
+ ),
218
+ (0, 1.5, Timestamp("20170103"), np.nan),
219
+ (
220
+ Timestamp("20170102", tz="US/Eastern"),
221
+ Timedelta("2 days"),
222
+ "baz",
223
+ pd.NaT,
224
+ ),
225
+ ],
226
+ )
227
+ def test_compare_list_like_object(self, op, interval_array, other):
228
+ result = op(interval_array, other)
229
+ expected = self.elementwise_comparison(op, interval_array, other)
230
+ tm.assert_numpy_array_equal(result, expected)
231
+
232
+ def test_compare_list_like_nan(self, op, interval_array, nulls_fixture):
233
+ other = [nulls_fixture] * 4
234
+ result = op(interval_array, other)
235
+ expected = self.elementwise_comparison(op, interval_array, other)
236
+
237
+ tm.assert_equal(result, expected)
238
+
239
+ @pytest.mark.parametrize(
240
+ "other",
241
+ [
242
+ np.arange(4, dtype="int64"),
243
+ np.arange(4, dtype="float64"),
244
+ date_range("2017-01-01", periods=4),
245
+ date_range("2017-01-01", periods=4, tz="US/Eastern"),
246
+ timedelta_range("0 days", periods=4),
247
+ period_range("2017-01-01", periods=4, freq="D"),
248
+ Categorical(list("abab")),
249
+ Categorical(date_range("2017-01-01", periods=4)),
250
+ pd.array(list("abcd")),
251
+ pd.array(["foo", 3.14, None, object()], dtype=object),
252
+ ],
253
+ ids=lambda x: str(x.dtype),
254
+ )
255
+ def test_compare_list_like_other(self, op, interval_array, other):
256
+ result = op(interval_array, other)
257
+ expected = self.elementwise_comparison(op, interval_array, other)
258
+ tm.assert_numpy_array_equal(result, expected)
259
+
260
+ @pytest.mark.parametrize("length", [1, 3, 5])
261
+ @pytest.mark.parametrize("other_constructor", [IntervalArray, list])
262
+ def test_compare_length_mismatch_errors(self, op, other_constructor, length):
263
+ interval_array = IntervalArray.from_arrays(range(4), range(1, 5))
264
+ other = other_constructor([Interval(0, 1)] * length)
265
+ with pytest.raises(ValueError, match="Lengths must match to compare"):
266
+ op(interval_array, other)
267
+
268
+ @pytest.mark.parametrize(
269
+ "constructor, expected_type, assert_func",
270
+ [
271
+ (IntervalIndex, np.array, tm.assert_numpy_array_equal),
272
+ (Series, Series, tm.assert_series_equal),
273
+ ],
274
+ )
275
+ def test_index_series_compat(self, op, constructor, expected_type, assert_func):
276
+ # IntervalIndex/Series that rely on IntervalArray for comparisons
277
+ breaks = range(4)
278
+ index = constructor(IntervalIndex.from_breaks(breaks))
279
+
280
+ # scalar comparisons
281
+ other = index[0]
282
+ result = op(index, other)
283
+ expected = expected_type(self.elementwise_comparison(op, index, other))
284
+ assert_func(result, expected)
285
+
286
+ other = breaks[0]
287
+ result = op(index, other)
288
+ expected = expected_type(self.elementwise_comparison(op, index, other))
289
+ assert_func(result, expected)
290
+
291
+ # list-like comparisons
292
+ other = IntervalArray.from_breaks(breaks)
293
+ result = op(index, other)
294
+ expected = expected_type(self.elementwise_comparison(op, index, other))
295
+ assert_func(result, expected)
296
+
297
+ other = [index[0], breaks[0], "foo"]
298
+ result = op(index, other)
299
+ expected = expected_type(self.elementwise_comparison(op, index, other))
300
+ assert_func(result, expected)
301
+
302
+ @pytest.mark.parametrize("scalars", ["a", False, 1, 1.0, None])
303
+ def test_comparison_operations(self, scalars):
304
+ # GH #28981
305
+ expected = Series([False, False])
306
+ s = Series([Interval(0, 1), Interval(1, 2)], dtype="interval")
307
+ result = s == scalars
308
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_numeric.py ADDED
@@ -0,0 +1,1588 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arithmetic tests for DataFrame/Series/Index/Array classes that should
2
+ # behave identically.
3
+ # Specifically for numeric dtypes
4
+ from __future__ import annotations
5
+
6
+ from collections import abc
7
+ from datetime import timedelta
8
+ from decimal import Decimal
9
+ import operator
10
+
11
+ import numpy as np
12
+ import pytest
13
+
14
+ import pandas as pd
15
+ from pandas import (
16
+ Index,
17
+ RangeIndex,
18
+ Series,
19
+ Timedelta,
20
+ TimedeltaIndex,
21
+ array,
22
+ date_range,
23
+ )
24
+ import pandas._testing as tm
25
+ from pandas.core import ops
26
+ from pandas.core.computation import expressions as expr
27
+ from pandas.tests.arithmetic.common import (
28
+ assert_invalid_addsub_type,
29
+ assert_invalid_comparison,
30
+ )
31
+
32
+
33
+ @pytest.fixture(autouse=True, params=[0, 1000000], ids=["numexpr", "python"])
34
+ def switch_numexpr_min_elements(request, monkeypatch):
35
+ with monkeypatch.context() as m:
36
+ m.setattr(expr, "_MIN_ELEMENTS", request.param)
37
+ yield request.param
38
+
39
+
40
+ @pytest.fixture(
41
+ params=[
42
+ # TODO: add more dtypes here
43
+ Index(np.arange(5, dtype="float64")),
44
+ Index(np.arange(5, dtype="int64")),
45
+ Index(np.arange(5, dtype="uint64")),
46
+ RangeIndex(5),
47
+ ],
48
+ ids=lambda x: type(x).__name__,
49
+ )
50
+ def numeric_idx(request):
51
+ """
52
+ Several types of numeric-dtypes Index objects
53
+ """
54
+ return request.param
55
+
56
+
57
+ def adjust_negative_zero(zero, expected):
58
+ """
59
+ Helper to adjust the expected result if we are dividing by -0.0
60
+ as opposed to 0.0
61
+ """
62
+ if np.signbit(np.array(zero)).any():
63
+ # All entries in the `zero` fixture should be either
64
+ # all-negative or no-negative.
65
+ assert np.signbit(np.array(zero)).all()
66
+
67
+ expected *= -1
68
+
69
+ return expected
70
+
71
+
72
+ def compare_op(series, other, op):
73
+ left = np.abs(series) if op in (ops.rpow, operator.pow) else series
74
+ right = np.abs(other) if op in (ops.rpow, operator.pow) else other
75
+
76
+ cython_or_numpy = op(left, right)
77
+ python = left.combine(right, op)
78
+ if isinstance(other, Series) and not other.index.equals(series.index):
79
+ python.index = python.index._with_freq(None)
80
+ tm.assert_series_equal(cython_or_numpy, python)
81
+
82
+
83
+ # TODO: remove this kludge once mypy stops giving false positives here
84
+ # List comprehension has incompatible type List[PandasObject]; expected List[RangeIndex]
85
+ # See GH#29725
86
+ _ldtypes = ["i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f2", "f4", "f8"]
87
+ lefts: list[Index | Series] = [RangeIndex(10, 40, 10)]
88
+ lefts.extend([Series([10, 20, 30], dtype=dtype) for dtype in _ldtypes])
89
+ lefts.extend([Index([10, 20, 30], dtype=dtype) for dtype in _ldtypes if dtype != "f2"])
90
+
91
+ # ------------------------------------------------------------------
92
+ # Comparisons
93
+
94
+
95
+ class TestNumericComparisons:
96
+ def test_operator_series_comparison_zerorank(self):
97
+ # GH#13006
98
+ result = np.float64(0) > Series([1, 2, 3])
99
+ expected = 0.0 > Series([1, 2, 3])
100
+ tm.assert_series_equal(result, expected)
101
+ result = Series([1, 2, 3]) < np.float64(0)
102
+ expected = Series([1, 2, 3]) < 0.0
103
+ tm.assert_series_equal(result, expected)
104
+ result = np.array([0, 1, 2])[0] > Series([0, 1, 2])
105
+ expected = 0.0 > Series([1, 2, 3])
106
+ tm.assert_series_equal(result, expected)
107
+
108
+ def test_df_numeric_cmp_dt64_raises(self, box_with_array, fixed_now_ts):
109
+ # GH#8932, GH#22163
110
+ ts = fixed_now_ts
111
+ obj = np.array(range(5))
112
+ obj = tm.box_expected(obj, box_with_array)
113
+
114
+ assert_invalid_comparison(obj, ts, box_with_array)
115
+
116
+ def test_compare_invalid(self):
117
+ # GH#8058
118
+ # ops testing
119
+ a = Series(np.random.default_rng(2).standard_normal(5), name=0)
120
+ b = Series(np.random.default_rng(2).standard_normal(5))
121
+ b.name = pd.Timestamp("2000-01-01")
122
+ tm.assert_series_equal(a / b, 1 / (b / a))
123
+
124
+ def test_numeric_cmp_string_numexpr_path(self, box_with_array, monkeypatch):
125
+ # GH#36377, GH#35700
126
+ box = box_with_array
127
+ xbox = box if box is not Index else np.ndarray
128
+
129
+ obj = Series(np.random.default_rng(2).standard_normal(51))
130
+ obj = tm.box_expected(obj, box, transpose=False)
131
+ with monkeypatch.context() as m:
132
+ m.setattr(expr, "_MIN_ELEMENTS", 50)
133
+ result = obj == "a"
134
+
135
+ expected = Series(np.zeros(51, dtype=bool))
136
+ expected = tm.box_expected(expected, xbox, transpose=False)
137
+ tm.assert_equal(result, expected)
138
+
139
+ with monkeypatch.context() as m:
140
+ m.setattr(expr, "_MIN_ELEMENTS", 50)
141
+ result = obj != "a"
142
+ tm.assert_equal(result, ~expected)
143
+
144
+ msg = "Invalid comparison between dtype=float64 and str"
145
+ with pytest.raises(TypeError, match=msg):
146
+ obj < "a"
147
+
148
+
149
+ # ------------------------------------------------------------------
150
+ # Numeric dtypes Arithmetic with Datetime/Timedelta Scalar
151
+
152
+
153
+ class TestNumericArraylikeArithmeticWithDatetimeLike:
154
+ def test_mul_timedelta_list(self, box_with_array):
155
+ # GH#62524
156
+ box = box_with_array
157
+ left = np.array([3, 4])
158
+ left = tm.box_expected(left, box)
159
+
160
+ right = [Timedelta(days=1), Timedelta(days=2)]
161
+
162
+ result = left * right
163
+
164
+ expected = TimedeltaIndex([Timedelta(days=3), Timedelta(days=8)])
165
+ expected = tm.box_expected(expected, box)
166
+ tm.assert_equal(result, expected)
167
+
168
+ result2 = right * left
169
+ tm.assert_equal(result2, expected)
170
+
171
+ @pytest.mark.parametrize("box_cls", [np.array, Index, Series])
172
+ @pytest.mark.parametrize(
173
+ "left", lefts, ids=lambda x: type(x).__name__ + str(x.dtype)
174
+ )
175
+ def test_mul_td64arr(self, left, box_cls):
176
+ # GH#22390
177
+ right = np.array([1, 2, 3], dtype="m8[s]")
178
+ right = box_cls(right)
179
+
180
+ expected = TimedeltaIndex(["10s", "40s", "90s"], dtype=right.dtype)
181
+
182
+ if isinstance(left, Series) or box_cls is Series:
183
+ expected = Series(expected)
184
+ assert expected.dtype == right.dtype
185
+
186
+ result = left * right
187
+ tm.assert_equal(result, expected)
188
+
189
+ result = right * left
190
+ tm.assert_equal(result, expected)
191
+
192
+ @pytest.mark.parametrize("box_cls", [np.array, Index, Series])
193
+ @pytest.mark.parametrize(
194
+ "left", lefts, ids=lambda x: type(x).__name__ + str(x.dtype)
195
+ )
196
+ def test_div_td64arr(self, left, box_cls):
197
+ # GH#22390
198
+ right = np.array([10, 40, 90], dtype="m8[s]")
199
+ right = box_cls(right)
200
+
201
+ expected = TimedeltaIndex(["1s", "2s", "3s"], dtype=right.dtype)
202
+ if isinstance(left, Series) or box_cls is Series:
203
+ expected = Series(expected)
204
+ assert expected.dtype == right.dtype
205
+
206
+ result = right / left
207
+ tm.assert_equal(result, expected)
208
+
209
+ result = right // left
210
+ tm.assert_equal(result, expected)
211
+
212
+ # (true_) needed for min-versions build 2022-12-26
213
+ msg = "ufunc '(true_)?divide' cannot use operands with types"
214
+ with pytest.raises(TypeError, match=msg):
215
+ left / right
216
+
217
+ msg = "ufunc 'floor_divide' cannot use operands with types"
218
+ with pytest.raises(TypeError, match=msg):
219
+ left // right
220
+
221
+ # TODO: also test Tick objects;
222
+ # see test_numeric_arr_rdiv_tdscalar for note on these failing
223
+ @pytest.mark.parametrize(
224
+ "scalar_td",
225
+ [
226
+ Timedelta(days=1).as_unit("ns"),
227
+ Timedelta(days=1).as_unit("ns").to_timedelta64(),
228
+ Timedelta(days=1).to_pytimedelta(),
229
+ Timedelta(days=1).to_timedelta64().astype("timedelta64[s]"),
230
+ Timedelta(days=1).to_timedelta64().astype("timedelta64[ms]"),
231
+ ],
232
+ ids=lambda x: type(x).__name__,
233
+ )
234
+ def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box_with_array):
235
+ # GH#19333
236
+ box = box_with_array
237
+ index = numeric_idx
238
+ expected = TimedeltaIndex(
239
+ [Timedelta(days=n) for n in range(len(index))], dtype="m8[ns]"
240
+ )
241
+ if isinstance(scalar_td, np.timedelta64):
242
+ dtype = scalar_td.dtype
243
+ expected = expected.astype(dtype)
244
+ elif type(scalar_td) is timedelta:
245
+ expected = expected.astype("m8[us]")
246
+
247
+ index = tm.box_expected(index, box)
248
+ expected = tm.box_expected(expected, box)
249
+
250
+ result = index * scalar_td
251
+ tm.assert_equal(result, expected)
252
+
253
+ commute = scalar_td * index
254
+ tm.assert_equal(commute, expected)
255
+
256
+ @pytest.mark.parametrize(
257
+ "scalar_td",
258
+ [
259
+ Timedelta(days=1).as_unit("ns"),
260
+ Timedelta(days=1).as_unit("ns").to_timedelta64(),
261
+ Timedelta(days=1).as_unit("ns").to_pytimedelta(),
262
+ ],
263
+ ids=lambda x: type(x).__name__,
264
+ )
265
+ @pytest.mark.parametrize("dtype", [np.int64, np.float64])
266
+ def test_numeric_arr_mul_tdscalar_numexpr_path(
267
+ self, dtype, scalar_td, box_with_array
268
+ ):
269
+ # GH#44772 for the float64 case
270
+ box = box_with_array
271
+
272
+ arr_i8 = np.arange(2 * 10**4).astype(np.int64, copy=False)
273
+ arr = arr_i8.astype(dtype, copy=False)
274
+ obj = tm.box_expected(arr, box, transpose=False)
275
+
276
+ expected = arr_i8.view("timedelta64[D]").astype("timedelta64[ns]")
277
+ if type(scalar_td) is timedelta:
278
+ expected = expected.astype("timedelta64[us]")
279
+
280
+ expected = tm.box_expected(expected, box, transpose=False)
281
+
282
+ result = obj * scalar_td
283
+ tm.assert_equal(result, expected)
284
+
285
+ result = scalar_td * obj
286
+ tm.assert_equal(result, expected)
287
+
288
+ def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box_with_array):
289
+ box = box_with_array
290
+
291
+ index = numeric_idx[1:3]
292
+
293
+ expected = TimedeltaIndex(["3 Days", "36 Hours"])
294
+ if isinstance(three_days, np.timedelta64):
295
+ dtype = three_days.dtype
296
+ if dtype < np.dtype("m8[s]"):
297
+ # i.e. resolution is lower -> use lowest supported resolution
298
+ dtype = np.dtype("m8[s]")
299
+ expected = expected.astype(dtype)
300
+ elif type(three_days) is timedelta or (
301
+ isinstance(three_days, Timedelta) and three_days.unit == "us"
302
+ ):
303
+ expected = expected.astype("m8[us]")
304
+ elif isinstance(
305
+ three_days,
306
+ (pd.offsets.Day, pd.offsets.Hour, pd.offsets.Minute, pd.offsets.Second),
307
+ ):
308
+ # closest reso is Second
309
+ expected = expected.astype("m8[s]")
310
+
311
+ index = tm.box_expected(index, box)
312
+ expected = tm.box_expected(expected, box)
313
+
314
+ if isinstance(three_days, pd.offsets.Day):
315
+ # GH#41943 Day is no longer timedelta-like
316
+ msg = "unsupported operand type"
317
+ with pytest.raises(TypeError, match=msg):
318
+ three_days / index
319
+ else:
320
+ result = three_days / index
321
+ tm.assert_equal(result, expected)
322
+ msg = "cannot use operands with types dtype"
323
+
324
+ with pytest.raises(TypeError, match=msg):
325
+ index / three_days
326
+
327
+ @pytest.mark.parametrize(
328
+ "other",
329
+ [
330
+ Timedelta(hours=31),
331
+ Timedelta(hours=31).to_pytimedelta(),
332
+ Timedelta(hours=31).to_timedelta64(),
333
+ Timedelta(hours=31).to_timedelta64().astype("m8[h]"),
334
+ np.timedelta64("NaT", "ns"),
335
+ np.timedelta64("NaT", "D"),
336
+ pd.offsets.Minute(3),
337
+ pd.offsets.Second(0),
338
+ # GH#28080 numeric+datetimelike should raise; Timestamp used
339
+ # to raise NullFrequencyError but that behavior was removed in 1.0
340
+ pd.Timestamp("2021-01-01", tz="Asia/Tokyo"),
341
+ pd.Timestamp("2021-01-01"),
342
+ pd.Timestamp("2021-01-01").to_pydatetime(),
343
+ pd.Timestamp("2021-01-01", tz="UTC").to_pydatetime(),
344
+ pd.Timestamp("2021-01-01").to_datetime64(),
345
+ np.datetime64("NaT", "ns"),
346
+ pd.NaT,
347
+ ],
348
+ ids=repr,
349
+ )
350
+ @pytest.mark.filterwarnings(
351
+ "ignore:.*'generic' unit for NumPy timedelta:DeprecationWarning"
352
+ )
353
+ def test_add_sub_datetimedeltalike_invalid(
354
+ self, numeric_idx, other, box_with_array
355
+ ):
356
+ box = box_with_array
357
+
358
+ left = tm.box_expected(numeric_idx, box)
359
+ msg = "|".join(
360
+ [
361
+ "unsupported operand type",
362
+ "Addition/subtraction of integers and integer-arrays",
363
+ "Instead of adding/subtracting",
364
+ "cannot use operands with types dtype",
365
+ "Concatenation operation is not implemented for NumPy arrays",
366
+ "Cannot (add|subtract) NaT (to|from) ndarray",
367
+ # pd.array vs np.datetime64 case
368
+ r"operand type\(s\) all returned NotImplemented from __array_ufunc__",
369
+ "can only perform ops with numeric values",
370
+ "cannot subtract DatetimeArray from ndarray",
371
+ # pd.Timedelta(1) + Index([0, 1, 2])
372
+ "Cannot add or subtract Timedelta from integers",
373
+ ]
374
+ )
375
+ assert_invalid_addsub_type(left, other, msg)
376
+
377
+
378
+ # ------------------------------------------------------------------
379
+ # Arithmetic
380
+
381
+
382
+ class TestDivisionByZero:
383
+ def test_div_zero(self, zero, numeric_idx):
384
+ idx = numeric_idx
385
+
386
+ expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
387
+ # We only adjust for Index, because Series does not yet apply
388
+ # the adjustment correctly.
389
+ expected2 = adjust_negative_zero(zero, expected)
390
+
391
+ result = idx / zero
392
+ tm.assert_index_equal(result, expected2)
393
+ ser_compat = Series(idx).astype("i8") / np.array(zero).astype("i8")
394
+ tm.assert_series_equal(ser_compat, Series(expected))
395
+
396
+ def test_floordiv_zero(self, zero, numeric_idx):
397
+ idx = numeric_idx
398
+
399
+ expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
400
+ # We only adjust for Index, because Series does not yet apply
401
+ # the adjustment correctly.
402
+ expected2 = adjust_negative_zero(zero, expected)
403
+
404
+ result = idx // zero
405
+ tm.assert_index_equal(result, expected2)
406
+ ser_compat = Series(idx).astype("i8") // np.array(zero).astype("i8")
407
+ tm.assert_series_equal(ser_compat, Series(expected))
408
+
409
+ def test_mod_zero(self, zero, numeric_idx):
410
+ idx = numeric_idx
411
+
412
+ expected = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=np.float64)
413
+ result = idx % zero
414
+ tm.assert_index_equal(result, expected)
415
+ ser_compat = Series(idx).astype("i8") % np.array(zero).astype("i8")
416
+ tm.assert_series_equal(ser_compat, Series(result))
417
+
418
+ def test_divmod_zero(self, zero, numeric_idx):
419
+ idx = numeric_idx
420
+
421
+ exleft = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
422
+ exright = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=np.float64)
423
+ exleft = adjust_negative_zero(zero, exleft)
424
+
425
+ result = divmod(idx, zero)
426
+ tm.assert_index_equal(result[0], exleft)
427
+ tm.assert_index_equal(result[1], exright)
428
+
429
+ @pytest.mark.parametrize("op", [operator.truediv, operator.floordiv])
430
+ def test_div_negative_zero(self, zero, numeric_idx, op):
431
+ # Check that -1 / -0.0 returns np.inf, not -np.inf
432
+ if numeric_idx.dtype == np.uint64:
433
+ pytest.skip(f"Div by negative 0 not relevant for {numeric_idx.dtype}")
434
+ idx = numeric_idx - 3
435
+
436
+ expected = Index([-np.inf, -np.inf, -np.inf, np.nan, np.inf], dtype=np.float64)
437
+ expected = adjust_negative_zero(zero, expected)
438
+
439
+ result = op(idx, zero)
440
+ tm.assert_index_equal(result, expected)
441
+
442
+ # ------------------------------------------------------------------
443
+
444
+ @pytest.mark.parametrize("dtype1", [np.int64, np.float64, np.uint64])
445
+ def test_ser_div_ser(
446
+ self,
447
+ switch_numexpr_min_elements,
448
+ dtype1,
449
+ any_real_numpy_dtype,
450
+ ):
451
+ # no longer do integer div for any ops, but deal with the 0's
452
+ dtype2 = any_real_numpy_dtype
453
+
454
+ first = Series([3, 4, 5, 8], name="first").astype(dtype1)
455
+ second = Series([0, 0, 0, 3], name="second").astype(dtype2)
456
+
457
+ with np.errstate(all="ignore"):
458
+ expected = Series(
459
+ first.values.astype(np.float64) / second.values,
460
+ dtype="float64",
461
+ name=None,
462
+ )
463
+ expected.iloc[0:3] = np.inf
464
+ if first.dtype == "int64" and second.dtype == "float32":
465
+ # when using numexpr, the casting rules are slightly different
466
+ # and int64/float32 combo results in float32 instead of float64
467
+ if expr.USE_NUMEXPR and switch_numexpr_min_elements == 0:
468
+ expected = expected.astype("float32")
469
+
470
+ result = first / second
471
+ tm.assert_series_equal(result, expected)
472
+ assert not result.equals(second / first)
473
+
474
+ @pytest.mark.parametrize("dtype1", [np.int64, np.float64, np.uint64])
475
+ def test_ser_divmod_zero(self, dtype1, any_real_numpy_dtype):
476
+ # GH#26987
477
+ dtype2 = any_real_numpy_dtype
478
+ left = Series([1, 1]).astype(dtype1)
479
+ right = Series([0, 2]).astype(dtype2)
480
+
481
+ # GH#27321 pandas convention is to set 1 // 0 to np.inf, as opposed
482
+ # to numpy which sets to np.nan; patch `expected[0]` below
483
+ expected = left // right, left % right
484
+ expected = list(expected)
485
+ expected[0] = expected[0].astype(np.float64)
486
+ expected[0][0] = np.inf
487
+ result = divmod(left, right)
488
+
489
+ tm.assert_series_equal(result[0], expected[0])
490
+ tm.assert_series_equal(result[1], expected[1])
491
+
492
+ # rdivmod case
493
+ result = divmod(left.values, right)
494
+ tm.assert_series_equal(result[0], expected[0])
495
+ tm.assert_series_equal(result[1], expected[1])
496
+
497
+ def test_ser_divmod_inf(self):
498
+ left = Series([np.inf, 1.0])
499
+ right = Series([np.inf, 2.0])
500
+
501
+ expected = left // right, left % right
502
+ result = divmod(left, right)
503
+
504
+ tm.assert_series_equal(result[0], expected[0])
505
+ tm.assert_series_equal(result[1], expected[1])
506
+
507
+ # rdivmod case
508
+ result = divmod(left.values, right)
509
+ tm.assert_series_equal(result[0], expected[0])
510
+ tm.assert_series_equal(result[1], expected[1])
511
+
512
+ def test_rdiv_zero_compat(self):
513
+ # GH#8674
514
+ zero_array = np.array([0] * 5)
515
+ data = np.random.default_rng(2).standard_normal(5)
516
+ expected = Series([0.0] * 5)
517
+
518
+ result = zero_array / Series(data)
519
+ tm.assert_series_equal(result, expected)
520
+
521
+ result = Series(zero_array) / data
522
+ tm.assert_series_equal(result, expected)
523
+
524
+ result = Series(zero_array) / Series(data)
525
+ tm.assert_series_equal(result, expected)
526
+
527
+ def test_div_zero_inf_signs(self):
528
+ # GH#9144, inf signing
529
+ ser = Series([-1, 0, 1], name="first")
530
+ expected = Series([-np.inf, np.nan, np.inf], name="first")
531
+
532
+ result = ser / 0
533
+ tm.assert_series_equal(result, expected)
534
+
535
+ def test_rdiv_zero(self):
536
+ # GH#9144
537
+ ser = Series([-1, 0, 1], name="first")
538
+ expected = Series([0.0, np.nan, 0.0], name="first")
539
+
540
+ result = 0 / ser
541
+ tm.assert_series_equal(result, expected)
542
+
543
+ def test_floordiv_div(self):
544
+ # GH#9144
545
+ ser = Series([-1, 0, 1], name="first")
546
+
547
+ result = ser // 0
548
+ expected = Series([-np.inf, np.nan, np.inf], name="first")
549
+ tm.assert_series_equal(result, expected)
550
+
551
+ def test_df_div_zero_df(self):
552
+ # integer div, but deal with the 0's (GH#9144)
553
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
554
+ result = df / df
555
+
556
+ first = Series([1.0, 1.0, 1.0, 1.0])
557
+ second = Series([np.nan, np.nan, np.nan, 1])
558
+ expected = pd.DataFrame({"first": first, "second": second})
559
+ tm.assert_frame_equal(result, expected)
560
+
561
+ def test_df_div_zero_array(self):
562
+ # integer div, but deal with the 0's (GH#9144)
563
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
564
+
565
+ first = Series([1.0, 1.0, 1.0, 1.0])
566
+ second = Series([np.nan, np.nan, np.nan, 1])
567
+ expected = pd.DataFrame({"first": first, "second": second})
568
+
569
+ with np.errstate(all="ignore"):
570
+ arr = df.values.astype("float") / df.values
571
+ result = pd.DataFrame(arr, index=df.index, columns=df.columns)
572
+ tm.assert_frame_equal(result, expected)
573
+
574
+ def test_df_div_zero_int(self):
575
+ # integer div, but deal with the 0's (GH#9144)
576
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
577
+
578
+ result = df / 0
579
+ expected = pd.DataFrame(np.inf, index=df.index, columns=df.columns)
580
+ expected.iloc[0:3, 1] = np.nan
581
+ tm.assert_frame_equal(result, expected)
582
+
583
+ # numpy has a slightly different (wrong) treatment
584
+ with np.errstate(all="ignore"):
585
+ arr = df.values.astype("float64") / 0
586
+ result2 = pd.DataFrame(arr, index=df.index, columns=df.columns)
587
+ tm.assert_frame_equal(result2, expected)
588
+
589
+ def test_df_div_zero_series_does_not_commute(self):
590
+ # integer div, but deal with the 0's (GH#9144)
591
+ df = pd.DataFrame(np.random.default_rng(2).standard_normal((10, 5)))
592
+ ser = df[0]
593
+ res = ser / df
594
+ res2 = df / ser
595
+ assert not res.fillna(0).equals(res2.fillna(0))
596
+
597
+ # ------------------------------------------------------------------
598
+ # Mod By Zero
599
+
600
+ def test_df_mod_zero_df(self):
601
+ # GH#3590, modulo as ints
602
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
603
+ # this is technically wrong, as the integer portion is coerced to float
604
+ first = Series([0, 0, 0, 0])
605
+ first = first.astype("float64")
606
+ second = Series([np.nan, np.nan, np.nan, 0])
607
+ expected = pd.DataFrame({"first": first, "second": second})
608
+ result = df % df
609
+ tm.assert_frame_equal(result, expected)
610
+
611
+ # GH#38939 If we dont pass copy=False, df is consolidated and
612
+ # result["first"] is float64 instead of int64
613
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]}, copy=False)
614
+ first = Series([0, 0, 0, 0], dtype="int64")
615
+ second = Series([np.nan, np.nan, np.nan, 0])
616
+ expected = pd.DataFrame({"first": first, "second": second})
617
+ result = df % df
618
+ tm.assert_frame_equal(result, expected)
619
+
620
+ def test_df_mod_zero_array(self):
621
+ # GH#3590, modulo as ints
622
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
623
+
624
+ # this is technically wrong, as the integer portion is coerced to float
625
+ # ###
626
+ first = Series([0, 0, 0, 0], dtype="float64")
627
+ second = Series([np.nan, np.nan, np.nan, 0])
628
+ expected = pd.DataFrame({"first": first, "second": second})
629
+
630
+ # numpy has a slightly different (wrong) treatment
631
+ with np.errstate(all="ignore"):
632
+ arr = df.values % df.values
633
+ result2 = pd.DataFrame(arr, index=df.index, columns=df.columns, dtype="float64")
634
+ result2.iloc[0:3, 1] = np.nan
635
+ tm.assert_frame_equal(result2, expected)
636
+
637
+ def test_df_mod_zero_int(self):
638
+ # GH#3590, modulo as ints
639
+ df = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
640
+
641
+ result = df % 0
642
+ expected = pd.DataFrame(np.nan, index=df.index, columns=df.columns)
643
+ tm.assert_frame_equal(result, expected)
644
+
645
+ # numpy has a slightly different (wrong) treatment
646
+ with np.errstate(all="ignore"):
647
+ arr = df.values.astype("float64") % 0
648
+ result2 = pd.DataFrame(arr, index=df.index, columns=df.columns)
649
+ tm.assert_frame_equal(result2, expected)
650
+
651
+ def test_df_mod_zero_series_does_not_commute(self):
652
+ # GH#3590, modulo as ints
653
+ # not commutative with series
654
+ df = pd.DataFrame(np.random.default_rng(2).standard_normal((10, 5)))
655
+ ser = df[0]
656
+ res = ser % df
657
+ res2 = df % ser
658
+ assert not res.fillna(0).equals(res2.fillna(0))
659
+
660
+
661
+ class TestMultiplicationDivision:
662
+ # __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__
663
+ # for non-timestamp/timedelta/period dtypes
664
+
665
+ def test_divide_decimal(self, box_with_array):
666
+ # resolves issue GH#9787
667
+ box = box_with_array
668
+ ser = Series([Decimal(10)])
669
+ expected = Series([Decimal(5)])
670
+
671
+ ser = tm.box_expected(ser, box)
672
+ expected = tm.box_expected(expected, box)
673
+
674
+ result = ser / Decimal(2)
675
+
676
+ tm.assert_equal(result, expected)
677
+
678
+ result = ser // Decimal(2)
679
+ tm.assert_equal(result, expected)
680
+
681
+ def test_div_equiv_binop(self):
682
+ # Test Series.div as well as Series.__div__
683
+ # float/integer issue
684
+ # GH#7785
685
+ first = Series([1, 0], name="first")
686
+ second = Series([-0.01, -0.02], name="second")
687
+ expected = Series([-0.01, -np.inf])
688
+
689
+ result = second.div(first)
690
+ tm.assert_series_equal(result, expected, check_names=False)
691
+
692
+ result = second / first
693
+ tm.assert_series_equal(result, expected)
694
+
695
+ def test_div_int(self, numeric_idx):
696
+ idx = numeric_idx
697
+ result = idx / 1
698
+ expected = idx.astype("float64")
699
+ tm.assert_index_equal(result, expected)
700
+
701
+ result = idx / 2
702
+ expected = Index(idx.values / 2)
703
+ tm.assert_index_equal(result, expected)
704
+
705
+ @pytest.mark.parametrize("op", [operator.mul, ops.rmul, operator.floordiv])
706
+ def test_mul_int_identity(self, op, numeric_idx, box_with_array):
707
+ idx = numeric_idx
708
+ idx = tm.box_expected(idx, box_with_array)
709
+
710
+ result = op(idx, 1)
711
+ tm.assert_equal(result, idx)
712
+
713
+ def test_mul_int_array(self, numeric_idx):
714
+ idx = numeric_idx
715
+ didx = idx * idx
716
+
717
+ result = idx * np.array(5, dtype="int64")
718
+ tm.assert_index_equal(result, idx * 5)
719
+
720
+ arr_dtype = "uint64" if idx.dtype == np.uint64 else "int64"
721
+ result = idx * np.arange(5, dtype=arr_dtype)
722
+ tm.assert_index_equal(result, didx)
723
+
724
+ def test_mul_int_series(self, numeric_idx):
725
+ idx = numeric_idx
726
+ didx = idx * idx
727
+
728
+ arr_dtype = "uint64" if idx.dtype == np.uint64 else "int64"
729
+ result = idx * Series(np.arange(5, dtype=arr_dtype))
730
+ tm.assert_series_equal(result, Series(didx))
731
+
732
+ def test_mul_float_series(self, numeric_idx):
733
+ idx = numeric_idx
734
+ rng5 = np.arange(5, dtype="float64")
735
+
736
+ result = idx * Series(rng5 + 0.1)
737
+ expected = Series(rng5 * (rng5 + 0.1))
738
+ tm.assert_series_equal(result, expected)
739
+
740
+ def test_mul_index(self, numeric_idx):
741
+ idx = numeric_idx
742
+
743
+ result = idx * idx
744
+ tm.assert_index_equal(result, idx**2)
745
+
746
+ def test_mul_datelike_raises(self, numeric_idx):
747
+ idx = numeric_idx
748
+ msg = "cannot perform __rmul__ with this index type"
749
+ with pytest.raises(TypeError, match=msg):
750
+ idx * date_range("20130101", periods=5)
751
+
752
+ def test_mul_size_mismatch_raises(self, numeric_idx):
753
+ idx = numeric_idx
754
+ msg = "operands could not be broadcast together"
755
+ with pytest.raises(ValueError, match=msg):
756
+ idx * idx[0:3]
757
+ with pytest.raises(ValueError, match=msg):
758
+ idx * np.array([1, 2])
759
+
760
+ @pytest.mark.parametrize("op", [operator.pow, ops.rpow])
761
+ def test_pow_float(self, op, numeric_idx, box_with_array):
762
+ # test power calculations both ways, GH#14973
763
+ box = box_with_array
764
+ idx = numeric_idx
765
+ expected = Index(op(idx.values, 2.0))
766
+
767
+ idx = tm.box_expected(idx, box)
768
+ expected = tm.box_expected(expected, box)
769
+
770
+ result = op(idx, 2.0)
771
+ tm.assert_equal(result, expected)
772
+
773
+ def test_modulo(self, numeric_idx, box_with_array):
774
+ # GH#9244
775
+ box = box_with_array
776
+ idx = numeric_idx
777
+ expected = Index(idx.values % 2)
778
+
779
+ idx = tm.box_expected(idx, box)
780
+ expected = tm.box_expected(expected, box)
781
+
782
+ result = idx % 2
783
+ tm.assert_equal(result, expected)
784
+
785
+ def test_divmod_scalar(self, numeric_idx):
786
+ idx = numeric_idx
787
+
788
+ result = divmod(idx, 2)
789
+ with np.errstate(all="ignore"):
790
+ div, mod = divmod(idx.values, 2)
791
+
792
+ expected = Index(div), Index(mod)
793
+ for r, e in zip(result, expected, strict=True):
794
+ tm.assert_index_equal(r, e)
795
+
796
+ def test_divmod_ndarray(self, numeric_idx):
797
+ idx = numeric_idx
798
+ other = np.ones(idx.values.shape, dtype=idx.values.dtype) * 2
799
+
800
+ result = divmod(idx, other)
801
+ with np.errstate(all="ignore"):
802
+ div, mod = divmod(idx.values, other)
803
+
804
+ expected = Index(div), Index(mod)
805
+ for r, e in zip(result, expected, strict=True):
806
+ tm.assert_index_equal(r, e)
807
+
808
+ def test_divmod_series(self, numeric_idx):
809
+ idx = numeric_idx
810
+ other = np.ones(idx.values.shape, dtype=idx.values.dtype) * 2
811
+
812
+ result = divmod(idx, Series(other))
813
+ with np.errstate(all="ignore"):
814
+ div, mod = divmod(idx.values, other)
815
+
816
+ expected = Series(div), Series(mod)
817
+ for r, e in zip(result, expected, strict=True):
818
+ tm.assert_series_equal(r, e)
819
+
820
+ @pytest.mark.parametrize("other", [np.nan, 7, -23, 2.718, -3.14, np.inf])
821
+ def test_ops_np_scalar(self, other):
822
+ vals = np.random.default_rng(2).standard_normal((5, 3))
823
+ f = lambda x: pd.DataFrame(
824
+ x, index=list("ABCDE"), columns=["jim", "joe", "jolie"]
825
+ )
826
+
827
+ df = f(vals)
828
+
829
+ tm.assert_frame_equal(df / np.array(other), f(vals / other))
830
+ tm.assert_frame_equal(np.array(other) * df, f(vals * other))
831
+ tm.assert_frame_equal(df + np.array(other), f(vals + other))
832
+ tm.assert_frame_equal(np.array(other) - df, f(other - vals))
833
+
834
+ # TODO: This came from series.test.test_operators, needs cleanup
835
+ def test_operators_frame(self):
836
+ # rpow does not work with DataFrame
837
+ ts = Series(
838
+ np.arange(10, dtype=np.float64),
839
+ index=date_range("2020-01-01", periods=10),
840
+ name="ts",
841
+ )
842
+ ts.name = "ts"
843
+
844
+ df = pd.DataFrame({"A": ts})
845
+
846
+ tm.assert_series_equal(ts + ts, ts + df["A"], check_names=False)
847
+ tm.assert_series_equal(ts**ts, ts ** df["A"], check_names=False)
848
+ tm.assert_series_equal(ts < ts, ts < df["A"], check_names=False)
849
+ tm.assert_series_equal(ts / ts, ts / df["A"], check_names=False)
850
+
851
+ # TODO: this came from tests.series.test_analytics, needs cleanup and
852
+ # de-duplication with test_modulo above
853
+ def test_modulo2(self):
854
+ with np.errstate(all="ignore"):
855
+ # GH#3590, modulo as ints
856
+ p = pd.DataFrame({"first": [3, 4, 5, 8], "second": [0, 0, 0, 3]})
857
+ result = p["first"] % p["second"]
858
+ expected = Series(p["first"].values % p["second"].values, dtype="float64")
859
+ expected.iloc[0:3] = np.nan
860
+ tm.assert_series_equal(result, expected)
861
+
862
+ result = p["first"] % 0
863
+ expected = Series(np.nan, index=p.index, name="first")
864
+ tm.assert_series_equal(result, expected)
865
+
866
+ p = p.astype("float64")
867
+ result = p["first"] % p["second"]
868
+ expected = Series(p["first"].values % p["second"].values)
869
+ tm.assert_series_equal(result, expected)
870
+
871
+ p = p.astype("float64")
872
+ result = p["first"] % p["second"]
873
+ result2 = p["second"] % p["first"]
874
+ assert not result.equals(result2)
875
+
876
+ def test_modulo_zero_int(self):
877
+ # GH#9144
878
+ with np.errstate(all="ignore"):
879
+ s = Series([0, 1])
880
+
881
+ result = s % 0
882
+ expected = Series([np.nan, np.nan])
883
+ tm.assert_series_equal(result, expected)
884
+
885
+ result = 0 % s
886
+ expected = Series([np.nan, 0.0])
887
+ tm.assert_series_equal(result, expected)
888
+
889
+ def test_non_1d_ea_raises_notimplementederror(self):
890
+ # GH#61866
891
+ ea_array = array([1, 2, 3, 4, 5], dtype="Int64").reshape(5, 1)
892
+ np_array = np.array([1, 2, 3, 4, 5], dtype=np.int64).reshape(5, 1)
893
+
894
+ msg = "can only perform ops with 1-d structures"
895
+
896
+ with pytest.raises(NotImplementedError, match=msg):
897
+ ea_array * np_array
898
+
899
+ with pytest.raises(NotImplementedError, match=msg):
900
+ np_array * ea_array
901
+
902
+
903
+ class TestAdditionSubtraction:
904
+ # __add__, __sub__, __radd__, __rsub__, __iadd__, __isub__
905
+ # for non-timestamp/timedelta/period dtypes
906
+
907
+ @pytest.mark.parametrize(
908
+ "first, second, expected",
909
+ [
910
+ (
911
+ Series([1, 2, 3], index=list("ABC"), name="x"),
912
+ Series([2, 2, 2], index=list("ABD"), name="x"),
913
+ Series([3.0, 4.0, np.nan, np.nan], index=list("ABCD"), name="x"),
914
+ ),
915
+ (
916
+ Series([1, 2, 3], index=list("ABC"), name="x"),
917
+ Series([2, 2, 2, 2], index=list("ABCD"), name="x"),
918
+ Series([3, 4, 5, np.nan], index=list("ABCD"), name="x"),
919
+ ),
920
+ ],
921
+ )
922
+ def test_add_series(self, first, second, expected):
923
+ # GH#1134
924
+ tm.assert_series_equal(first + second, expected)
925
+ tm.assert_series_equal(second + first, expected)
926
+
927
+ @pytest.mark.parametrize(
928
+ "first, second, expected",
929
+ [
930
+ (
931
+ pd.DataFrame({"x": [1, 2, 3]}, index=list("ABC")),
932
+ pd.DataFrame({"x": [2, 2, 2]}, index=list("ABD")),
933
+ pd.DataFrame({"x": [3.0, 4.0, np.nan, np.nan]}, index=list("ABCD")),
934
+ ),
935
+ (
936
+ pd.DataFrame({"x": [1, 2, 3]}, index=list("ABC")),
937
+ pd.DataFrame({"x": [2, 2, 2, 2]}, index=list("ABCD")),
938
+ pd.DataFrame({"x": [3, 4, 5, np.nan]}, index=list("ABCD")),
939
+ ),
940
+ ],
941
+ )
942
+ def test_add_frames(self, first, second, expected):
943
+ # GH#1134
944
+ tm.assert_frame_equal(first + second, expected)
945
+ tm.assert_frame_equal(second + first, expected)
946
+
947
+ # TODO: This came from series.test.test_operators, needs cleanup
948
+ def test_series_frame_radd_bug(self, fixed_now_ts):
949
+ # GH#353
950
+ vals = Series([str(i) for i in range(5)])
951
+ result = "foo_" + vals
952
+ expected = vals.map(lambda x: "foo_" + x)
953
+ tm.assert_series_equal(result, expected)
954
+
955
+ frame = pd.DataFrame({"vals": vals})
956
+ result = "foo_" + frame
957
+ expected = pd.DataFrame({"vals": vals.map(lambda x: "foo_" + x)})
958
+ tm.assert_frame_equal(result, expected)
959
+
960
+ ts = Series(
961
+ np.arange(10, dtype=np.float64),
962
+ index=date_range("2020-01-01", periods=10),
963
+ name="ts",
964
+ )
965
+
966
+ # really raise this time
967
+ fix_now = fixed_now_ts.to_pydatetime()
968
+ msg = "|".join(
969
+ [
970
+ "unsupported operand type",
971
+ # wrong error message, see https://github.com/numpy/numpy/issues/18832
972
+ "Concatenation operation",
973
+ ]
974
+ )
975
+ with pytest.raises(TypeError, match=msg):
976
+ fix_now + ts
977
+
978
+ with pytest.raises(TypeError, match=msg):
979
+ ts + fix_now
980
+
981
+ # TODO: This came from series.test.test_operators, needs cleanup
982
+ def test_datetime64_with_index(self):
983
+ # arithmetic integer ops with an index
984
+ ser = Series(np.random.default_rng(2).standard_normal(5))
985
+ expected = ser - ser.index.to_series()
986
+ result = ser - ser.index
987
+ tm.assert_series_equal(result, expected)
988
+
989
+ # GH#4629
990
+ # arithmetic datetime64 ops with an index
991
+ ser = Series(
992
+ date_range("20130101", periods=5),
993
+ index=date_range("20130101", periods=5),
994
+ )
995
+ expected = ser - ser.index.to_series()
996
+ result = ser - ser.index
997
+ tm.assert_series_equal(result, expected)
998
+
999
+ msg = "cannot subtract PeriodArray from DatetimeArray"
1000
+ with pytest.raises(TypeError, match=msg):
1001
+ # GH#18850
1002
+ result = ser - ser.index.to_period()
1003
+
1004
+ df = pd.DataFrame(
1005
+ np.random.default_rng(2).standard_normal((5, 2)),
1006
+ index=date_range("20130101", periods=5),
1007
+ )
1008
+ df["date"] = pd.Timestamp("20130102")
1009
+ df["expected"] = df["date"] - df.index.to_series()
1010
+ df["result"] = df["date"] - df.index
1011
+ tm.assert_series_equal(df["result"], df["expected"], check_names=False)
1012
+
1013
+ # TODO: taken from tests.frame.test_operators, needs cleanup
1014
+ def test_frame_operators(self, float_frame):
1015
+ frame = float_frame
1016
+
1017
+ garbage = np.random.default_rng(2).random(4)
1018
+ colSeries = Series(garbage, index=np.array(frame.columns))
1019
+
1020
+ idSum = frame + frame
1021
+ seriesSum = frame + colSeries
1022
+
1023
+ for col, series in idSum.items():
1024
+ for idx, val in series.items():
1025
+ origVal = frame[col][idx] * 2
1026
+ if not np.isnan(val):
1027
+ assert val == origVal
1028
+ else:
1029
+ assert np.isnan(origVal)
1030
+
1031
+ for col, series in seriesSum.items():
1032
+ for idx, val in series.items():
1033
+ origVal = frame[col][idx] + colSeries[col]
1034
+ if not np.isnan(val):
1035
+ assert val == origVal
1036
+ else:
1037
+ assert np.isnan(origVal)
1038
+
1039
+ def test_frame_operators_col_align(self, float_frame):
1040
+ frame2 = pd.DataFrame(float_frame, columns=["D", "C", "B", "A"])
1041
+ added = frame2 + frame2
1042
+ expected = frame2 * 2
1043
+ tm.assert_frame_equal(added, expected)
1044
+
1045
+ def test_frame_operators_none_to_nan(self):
1046
+ df = pd.DataFrame({"a": ["a", None, "b"]})
1047
+ tm.assert_frame_equal(df + df, pd.DataFrame({"a": ["aa", np.nan, "bb"]}))
1048
+
1049
+ @pytest.mark.parametrize("dtype", ("float", "int64"))
1050
+ def test_frame_operators_empty_like(self, dtype):
1051
+ # Test for issue #10181
1052
+ frames = [
1053
+ pd.DataFrame(dtype=dtype),
1054
+ pd.DataFrame(columns=["A"], dtype=dtype),
1055
+ pd.DataFrame(index=[0], dtype=dtype),
1056
+ ]
1057
+ for df in frames:
1058
+ assert (df + df).equals(df)
1059
+ tm.assert_frame_equal(df + df, df)
1060
+
1061
+ @pytest.mark.parametrize(
1062
+ "func",
1063
+ [lambda x: x * 2, lambda x: x[::2], lambda x: 5],
1064
+ ids=["multiply", "slice", "constant"],
1065
+ )
1066
+ def test_series_operators_arithmetic(self, all_arithmetic_functions, func):
1067
+ op = all_arithmetic_functions
1068
+ series = Series(
1069
+ np.arange(10, dtype=np.float64),
1070
+ index=date_range("2020-01-01", periods=10),
1071
+ name="ts",
1072
+ )
1073
+ other = func(series)
1074
+ compare_op(series, other, op)
1075
+
1076
+ @pytest.mark.parametrize(
1077
+ "func", [lambda x: x + 1, lambda x: 5], ids=["add", "constant"]
1078
+ )
1079
+ def test_series_operators_compare(self, comparison_op, func):
1080
+ op = comparison_op
1081
+ series = Series(
1082
+ np.arange(10, dtype=np.float64),
1083
+ index=date_range("2020-01-01", periods=10),
1084
+ name="ts",
1085
+ )
1086
+ other = func(series)
1087
+ compare_op(series, other, op)
1088
+
1089
+ @pytest.mark.parametrize(
1090
+ "func",
1091
+ [lambda x: x * 2, lambda x: x[::2], lambda x: 5],
1092
+ ids=["multiply", "slice", "constant"],
1093
+ )
1094
+ def test_divmod(self, func):
1095
+ series = Series(
1096
+ np.arange(10, dtype=np.float64),
1097
+ index=date_range("2020-01-01", periods=10),
1098
+ name="ts",
1099
+ )
1100
+ other = func(series)
1101
+ results = divmod(series, other)
1102
+ if isinstance(other, abc.Iterable) and len(series) != len(other):
1103
+ # if the lengths don't match, this is the test where we use
1104
+ # `tser[::2]`. Pad every other value in `other_np` with nan.
1105
+ other_np = []
1106
+ for n in other:
1107
+ other_np.append(n)
1108
+ other_np.append(np.nan)
1109
+ else:
1110
+ other_np = other
1111
+ other_np = np.asarray(other_np)
1112
+ with np.errstate(all="ignore"):
1113
+ expecteds = divmod(series.values, np.asarray(other_np))
1114
+
1115
+ for result, expected in zip(results, expecteds, strict=True):
1116
+ # check the values, name, and index separately
1117
+ tm.assert_almost_equal(np.asarray(result), expected)
1118
+
1119
+ assert result.name == series.name
1120
+ tm.assert_index_equal(result.index, series.index._with_freq(None))
1121
+
1122
+ def test_series_divmod_zero(self):
1123
+ # Check that divmod uses pandas convention for division by zero,
1124
+ # which does not match numpy.
1125
+ # pandas convention has
1126
+ # 1/0 == np.inf
1127
+ # -1/0 == -np.inf
1128
+ # 1/-0.0 == -np.inf
1129
+ # -1/-0.0 == np.inf
1130
+ tser = Series(
1131
+ np.arange(1, 11, dtype=np.float64),
1132
+ index=date_range("2020-01-01", periods=10),
1133
+ name="ts",
1134
+ )
1135
+ other = tser * 0
1136
+
1137
+ result = divmod(tser, other)
1138
+ exp1 = Series([np.inf] * len(tser), index=tser.index, name="ts")
1139
+ exp2 = Series([np.nan] * len(tser), index=tser.index, name="ts")
1140
+ tm.assert_series_equal(result[0], exp1)
1141
+ tm.assert_series_equal(result[1], exp2)
1142
+
1143
+
1144
+ class TestUFuncCompat:
1145
+ # TODO: add more dtypes
1146
+ @pytest.mark.parametrize("holder", [Index, RangeIndex, Series])
1147
+ @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64])
1148
+ def test_ufunc_compat(self, holder, dtype):
1149
+ box = Series if holder is Series else Index
1150
+
1151
+ if holder is RangeIndex:
1152
+ if dtype != np.int64:
1153
+ pytest.skip(f"dtype {dtype} not relevant for RangeIndex")
1154
+ idx = RangeIndex(0, 5, name="foo")
1155
+ else:
1156
+ idx = holder(np.arange(5, dtype=dtype), name="foo")
1157
+ result = np.sin(idx)
1158
+ expected = box(np.sin(np.arange(5, dtype=dtype)), name="foo")
1159
+ tm.assert_equal(result, expected)
1160
+
1161
+ # TODO: add more dtypes
1162
+ @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64])
1163
+ def test_ufunc_coercions(self, index_or_series, dtype):
1164
+ idx = index_or_series([1, 2, 3, 4, 5], dtype=dtype, name="x")
1165
+ box = index_or_series
1166
+
1167
+ result = np.sqrt(idx)
1168
+ assert result.dtype == "f8" and isinstance(result, box)
1169
+ exp = Index(np.sqrt(np.array([1, 2, 3, 4, 5], dtype=np.float64)), name="x")
1170
+ exp = tm.box_expected(exp, box)
1171
+ tm.assert_equal(result, exp)
1172
+
1173
+ result = np.divide(idx, 2.0)
1174
+ assert result.dtype == "f8" and isinstance(result, box)
1175
+ exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=np.float64, name="x")
1176
+ exp = tm.box_expected(exp, box)
1177
+ tm.assert_equal(result, exp)
1178
+
1179
+ # _evaluate_numeric_binop
1180
+ result = idx + 2.0
1181
+ assert result.dtype == "f8" and isinstance(result, box)
1182
+ exp = Index([3.0, 4.0, 5.0, 6.0, 7.0], dtype=np.float64, name="x")
1183
+ exp = tm.box_expected(exp, box)
1184
+ tm.assert_equal(result, exp)
1185
+
1186
+ result = idx - 2.0
1187
+ assert result.dtype == "f8" and isinstance(result, box)
1188
+ exp = Index([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=np.float64, name="x")
1189
+ exp = tm.box_expected(exp, box)
1190
+ tm.assert_equal(result, exp)
1191
+
1192
+ result = idx * 1.0
1193
+ assert result.dtype == "f8" and isinstance(result, box)
1194
+ exp = Index([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float64, name="x")
1195
+ exp = tm.box_expected(exp, box)
1196
+ tm.assert_equal(result, exp)
1197
+
1198
+ result = idx / 2.0
1199
+ assert result.dtype == "f8" and isinstance(result, box)
1200
+ exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=np.float64, name="x")
1201
+ exp = tm.box_expected(exp, box)
1202
+ tm.assert_equal(result, exp)
1203
+
1204
+ # TODO: add more dtypes
1205
+ @pytest.mark.parametrize("holder", [Index, Series])
1206
+ @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64])
1207
+ def test_ufunc_multiple_return_values(self, holder, dtype):
1208
+ obj = holder([1, 2, 3], dtype=dtype, name="x")
1209
+ box = Series if holder is Series else Index
1210
+
1211
+ result = np.modf(obj)
1212
+ assert isinstance(result, tuple)
1213
+ exp1 = Index([0.0, 0.0, 0.0], dtype=np.float64, name="x")
1214
+ exp2 = Index([1.0, 2.0, 3.0], dtype=np.float64, name="x")
1215
+ tm.assert_equal(result[0], tm.box_expected(exp1, box))
1216
+ tm.assert_equal(result[1], tm.box_expected(exp2, box))
1217
+
1218
+ def test_ufunc_at(self):
1219
+ s = Series([0, 1, 2], index=[1, 2, 3], name="x")
1220
+ np.add.at(s, [0, 2], 10)
1221
+ expected = Series([10, 1, 12], index=[1, 2, 3], name="x")
1222
+ tm.assert_series_equal(s, expected)
1223
+
1224
+
1225
+ class TestObjectDtypeEquivalence:
1226
+ # Tests that arithmetic operations match operations executed elementwise
1227
+
1228
+ @pytest.mark.parametrize("dtype", [None, object])
1229
+ def test_numarr_with_dtype_add_nan(self, dtype, box_with_array):
1230
+ box = box_with_array
1231
+ ser = Series([1, 2, 3], dtype=dtype)
1232
+ expected = Series([np.nan, np.nan, np.nan], dtype=dtype)
1233
+
1234
+ ser = tm.box_expected(ser, box)
1235
+ expected = tm.box_expected(expected, box)
1236
+
1237
+ result = np.nan + ser
1238
+ tm.assert_equal(result, expected)
1239
+
1240
+ result = ser + np.nan
1241
+ tm.assert_equal(result, expected)
1242
+
1243
+ @pytest.mark.parametrize("dtype", [None, object])
1244
+ def test_numarr_with_dtype_add_int(self, dtype, box_with_array):
1245
+ box = box_with_array
1246
+ ser = Series([1, 2, 3], dtype=dtype)
1247
+ expected = Series([2, 3, 4], dtype=dtype)
1248
+
1249
+ ser = tm.box_expected(ser, box)
1250
+ expected = tm.box_expected(expected, box)
1251
+
1252
+ result = 1 + ser
1253
+ tm.assert_equal(result, expected)
1254
+
1255
+ result = ser + 1
1256
+ tm.assert_equal(result, expected)
1257
+
1258
+ # TODO: moved from tests.series.test_operators; needs cleanup
1259
+ @pytest.mark.parametrize(
1260
+ "op",
1261
+ [operator.add, operator.sub, operator.mul, operator.truediv, operator.floordiv],
1262
+ )
1263
+ def test_operators_reverse_object(self, op):
1264
+ # GH#56
1265
+ arr = Series(
1266
+ np.random.default_rng(2).standard_normal(10),
1267
+ index=np.arange(10),
1268
+ dtype=object,
1269
+ )
1270
+
1271
+ result = op(1.0, arr)
1272
+ expected = op(1.0, arr.astype(float))
1273
+ tm.assert_series_equal(result.astype(float), expected)
1274
+
1275
+
1276
+ class TestNumericArithmeticUnsorted:
1277
+ # Tests in this class have been moved from type-specific test modules
1278
+ # but not yet sorted, parametrized, and de-duplicated
1279
+ @pytest.mark.parametrize(
1280
+ "op",
1281
+ [
1282
+ operator.add,
1283
+ operator.sub,
1284
+ operator.mul,
1285
+ operator.floordiv,
1286
+ operator.truediv,
1287
+ ],
1288
+ )
1289
+ @pytest.mark.parametrize(
1290
+ "idx1",
1291
+ [
1292
+ RangeIndex(0, 10, 1),
1293
+ RangeIndex(0, 20, 2),
1294
+ RangeIndex(-10, 10, 2),
1295
+ RangeIndex(5, -5, -1),
1296
+ ],
1297
+ )
1298
+ @pytest.mark.parametrize(
1299
+ "idx2",
1300
+ [
1301
+ RangeIndex(0, 10, 1),
1302
+ RangeIndex(0, 20, 2),
1303
+ RangeIndex(-10, 10, 2),
1304
+ RangeIndex(5, -5, -1),
1305
+ ],
1306
+ )
1307
+ def test_binops_index(self, op, idx1, idx2):
1308
+ idx1 = idx1._rename("foo")
1309
+ idx2 = idx2._rename("bar")
1310
+ result = op(idx1, idx2)
1311
+ expected = op(Index(idx1.to_numpy()), Index(idx2.to_numpy()))
1312
+ tm.assert_index_equal(result, expected, exact="equiv")
1313
+
1314
+ @pytest.mark.parametrize(
1315
+ "op",
1316
+ [
1317
+ operator.add,
1318
+ operator.sub,
1319
+ operator.mul,
1320
+ operator.floordiv,
1321
+ operator.truediv,
1322
+ ],
1323
+ )
1324
+ @pytest.mark.parametrize(
1325
+ "idx",
1326
+ [
1327
+ RangeIndex(0, 10, 1),
1328
+ RangeIndex(0, 20, 2),
1329
+ RangeIndex(-10, 10, 2),
1330
+ RangeIndex(5, -5, -1),
1331
+ ],
1332
+ )
1333
+ @pytest.mark.parametrize("scalar", [-1, 1, 2])
1334
+ def test_binops_index_scalar(self, op, idx, scalar):
1335
+ result = op(idx, scalar)
1336
+ expected = op(Index(idx.to_numpy()), scalar)
1337
+ tm.assert_index_equal(result, expected, exact="equiv")
1338
+
1339
+ @pytest.mark.parametrize("idx1", [RangeIndex(0, 10, 1), RangeIndex(0, 20, 2)])
1340
+ @pytest.mark.parametrize("idx2", [RangeIndex(0, 10, 1), RangeIndex(0, 20, 2)])
1341
+ def test_binops_index_pow(self, idx1, idx2):
1342
+ # numpy does not allow powers of negative integers so test separately
1343
+ # https://github.com/numpy/numpy/pull/8127
1344
+ idx1 = idx1._rename("foo")
1345
+ idx2 = idx2._rename("bar")
1346
+ result = pow(idx1, idx2)
1347
+ expected = pow(Index(idx1.to_numpy()), Index(idx2.to_numpy()))
1348
+ tm.assert_index_equal(result, expected, exact="equiv")
1349
+
1350
+ @pytest.mark.parametrize("idx", [RangeIndex(0, 10, 1), RangeIndex(0, 20, 2)])
1351
+ @pytest.mark.parametrize("scalar", [1, 2])
1352
+ def test_binops_index_scalar_pow(self, idx, scalar):
1353
+ # numpy does not allow powers of negative integers so test separately
1354
+ # https://github.com/numpy/numpy/pull/8127
1355
+ result = pow(idx, scalar)
1356
+ expected = pow(Index(idx.to_numpy()), scalar)
1357
+ tm.assert_index_equal(result, expected, exact="equiv")
1358
+
1359
+ # TODO: divmod?
1360
+ @pytest.mark.parametrize(
1361
+ "op",
1362
+ [
1363
+ operator.add,
1364
+ operator.sub,
1365
+ operator.mul,
1366
+ operator.floordiv,
1367
+ operator.truediv,
1368
+ operator.pow,
1369
+ operator.mod,
1370
+ ],
1371
+ )
1372
+ def test_arithmetic_with_frame_or_series(self, op):
1373
+ # check that we return NotImplemented when operating with Series
1374
+ # or DataFrame
1375
+ index = RangeIndex(5)
1376
+ other = Series(np.random.default_rng(2).standard_normal(5))
1377
+
1378
+ expected = op(Series(index), other)
1379
+ result = op(index, other)
1380
+ tm.assert_series_equal(result, expected)
1381
+
1382
+ other = pd.DataFrame(np.random.default_rng(2).standard_normal((2, 5)))
1383
+ expected = op(pd.DataFrame([index, index]), other)
1384
+ result = op(index, other)
1385
+ tm.assert_frame_equal(result, expected)
1386
+
1387
+ def test_numeric_compat2(self):
1388
+ # validate that we are handling the RangeIndex overrides to numeric ops
1389
+ # and returning RangeIndex where possible
1390
+
1391
+ idx = RangeIndex(0, 10, 2)
1392
+
1393
+ result = idx * 2
1394
+ expected = RangeIndex(0, 20, 4)
1395
+ tm.assert_index_equal(result, expected, exact=True)
1396
+
1397
+ result = idx + 2
1398
+ expected = RangeIndex(2, 12, 2)
1399
+ tm.assert_index_equal(result, expected, exact=True)
1400
+
1401
+ result = idx - 2
1402
+ expected = RangeIndex(-2, 8, 2)
1403
+ tm.assert_index_equal(result, expected, exact=True)
1404
+
1405
+ result = idx / 2
1406
+ expected = RangeIndex(0, 5, 1).astype("float64")
1407
+ tm.assert_index_equal(result, expected, exact=True)
1408
+
1409
+ result = idx / 4
1410
+ expected = RangeIndex(0, 10, 2) / 4
1411
+ tm.assert_index_equal(result, expected, exact=True)
1412
+
1413
+ result = idx // 1
1414
+ expected = idx
1415
+ tm.assert_index_equal(result, expected, exact=True)
1416
+
1417
+ # __mul__
1418
+ result = idx * idx
1419
+ expected = Index(idx.values * idx.values)
1420
+ tm.assert_index_equal(result, expected, exact=True)
1421
+
1422
+ # __pow__
1423
+ idx = RangeIndex(0, 1000, 2)
1424
+ result = idx**2
1425
+ expected = Index(idx._values) ** 2
1426
+ tm.assert_index_equal(Index(result.values), expected, exact=True)
1427
+
1428
+ @pytest.mark.parametrize(
1429
+ "idx, div, expected",
1430
+ [
1431
+ # TODO: add more dtypes
1432
+ (RangeIndex(0, 1000, 2), 2, RangeIndex(0, 500, 1)),
1433
+ (RangeIndex(-99, -201, -3), -3, RangeIndex(33, 67, 1)),
1434
+ (
1435
+ RangeIndex(0, 1000, 1),
1436
+ 2,
1437
+ Index(RangeIndex(0, 1000, 1)._values) // 2,
1438
+ ),
1439
+ (
1440
+ RangeIndex(0, 100, 1),
1441
+ 2.0,
1442
+ Index(RangeIndex(0, 100, 1)._values) // 2.0,
1443
+ ),
1444
+ (RangeIndex(0), 50, RangeIndex(0)),
1445
+ (RangeIndex(2, 4, 2), 3, RangeIndex(0, 1, 1)),
1446
+ (RangeIndex(-5, -10, -6), 4, RangeIndex(-2, -1, 1)),
1447
+ (RangeIndex(-100, -200, 3), 2, RangeIndex(0)),
1448
+ ],
1449
+ )
1450
+ def test_numeric_compat2_floordiv(self, idx, div, expected):
1451
+ # __floordiv__
1452
+ tm.assert_index_equal(idx // div, expected, exact=True)
1453
+
1454
+ @pytest.mark.parametrize("dtype", [np.int64, np.float64])
1455
+ @pytest.mark.parametrize("delta", [1, 0, -1])
1456
+ def test_addsub_arithmetic(self, dtype, delta):
1457
+ # GH#8142
1458
+ delta = dtype(delta)
1459
+ index = Index([10, 11, 12], dtype=dtype)
1460
+ result = index + delta
1461
+ expected = Index(index.values + delta, dtype=dtype)
1462
+ tm.assert_index_equal(result, expected)
1463
+
1464
+ # this subtraction used to fail
1465
+ result = index - delta
1466
+ expected = Index(index.values - delta, dtype=dtype)
1467
+ tm.assert_index_equal(result, expected)
1468
+
1469
+ tm.assert_index_equal(index + index, 2 * index)
1470
+ tm.assert_index_equal(index - index, 0 * index)
1471
+ assert not (index - index).empty
1472
+
1473
+ def test_pow_nan_with_zero(self, box_with_array):
1474
+ left = Index([np.nan, np.nan, np.nan])
1475
+ right = Index([0, 0, 0])
1476
+ expected = Index([1.0, 1.0, 1.0])
1477
+
1478
+ left = tm.box_expected(left, box_with_array)
1479
+ right = tm.box_expected(right, box_with_array)
1480
+ expected = tm.box_expected(expected, box_with_array)
1481
+
1482
+ result = left**right
1483
+ tm.assert_equal(result, expected)
1484
+
1485
+
1486
+ def test_fill_value_inf_masking():
1487
+ # GH #27464 make sure we mask 0/1 with Inf and not NaN
1488
+ df = pd.DataFrame({"A": [0, 1, 2], "B": [1.1, None, 1.1]})
1489
+
1490
+ other = pd.DataFrame({"A": [1.1, 1.2, 1.3]}, index=[0, 2, 3])
1491
+
1492
+ result = df.rfloordiv(other, fill_value=1)
1493
+
1494
+ expected = pd.DataFrame(
1495
+ {"A": [np.inf, 1.0, 0.0, 1.0], "B": [0.0, np.nan, 0.0, np.nan]}
1496
+ )
1497
+ tm.assert_frame_equal(result, expected, check_index_type=False)
1498
+
1499
+
1500
+ def test_dataframe_div_silenced():
1501
+ # GH#26793
1502
+ pdf1 = pd.DataFrame(
1503
+ {
1504
+ "A": np.arange(10),
1505
+ "B": [np.nan, 1, 2, 3, 4] * 2,
1506
+ "C": [np.nan] * 10,
1507
+ "D": np.arange(10),
1508
+ },
1509
+ index=list("abcdefghij"),
1510
+ columns=list("ABCD"),
1511
+ )
1512
+ pdf2 = pd.DataFrame(
1513
+ np.random.default_rng(2).standard_normal((10, 4)),
1514
+ index=list("abcdefghjk"),
1515
+ columns=list("ABCX"),
1516
+ )
1517
+ with tm.assert_produces_warning(None):
1518
+ pdf1.div(pdf2, fill_value=0)
1519
+
1520
+
1521
+ @pytest.mark.parametrize(
1522
+ "data, expected_data",
1523
+ [([0, 1, 2], [0, 2, 4])],
1524
+ )
1525
+ @pytest.mark.parametrize("box_pandas_1d_array", [Index, Series, tm.to_array])
1526
+ @pytest.mark.parametrize("box_1d_array", [Index, Series, tm.to_array, np.array, list])
1527
+ def test_integer_array_add_list_like(
1528
+ box_pandas_1d_array, box_1d_array, data, expected_data
1529
+ ):
1530
+ # GH22606 Verify operators with IntegerArray and list-likes
1531
+ arr = array(data, dtype="Int64")
1532
+ container = box_pandas_1d_array(arr)
1533
+ left = container + box_1d_array(data)
1534
+ right = box_1d_array(data) + container
1535
+
1536
+ if Series in [box_1d_array, box_pandas_1d_array]:
1537
+ cls = Series
1538
+ elif Index in [box_1d_array, box_pandas_1d_array]:
1539
+ cls = Index
1540
+ else:
1541
+ cls = array
1542
+
1543
+ expected = cls(expected_data, dtype="Int64")
1544
+
1545
+ tm.assert_equal(left, expected)
1546
+ tm.assert_equal(right, expected)
1547
+
1548
+
1549
+ def test_sub_multiindex_swapped_levels():
1550
+ # GH 9952
1551
+ df = pd.DataFrame(
1552
+ {"a": np.random.default_rng(2).standard_normal(6)},
1553
+ index=pd.MultiIndex.from_product(
1554
+ [["a", "b"], [0, 1, 2]], names=["levA", "levB"]
1555
+ ),
1556
+ )
1557
+ df2 = df.copy()
1558
+ df2.index = df2.index.swaplevel(0, 1)
1559
+ result = df - df2
1560
+ expected = pd.DataFrame([0.0] * 6, columns=["a"], index=df.index)
1561
+ tm.assert_frame_equal(result, expected)
1562
+
1563
+
1564
+ @pytest.mark.parametrize("power", [1, 2, 5])
1565
+ @pytest.mark.parametrize("string_size", [0, 1, 2, 5])
1566
+ def test_empty_str_comparison(power, string_size):
1567
+ # GH 37348
1568
+ a = np.array(range(10**power))
1569
+ right = pd.DataFrame(a, dtype=np.int64)
1570
+ left = " " * string_size
1571
+
1572
+ result = right == left
1573
+ expected = pd.DataFrame(np.zeros(right.shape, dtype=bool))
1574
+ tm.assert_frame_equal(result, expected)
1575
+
1576
+
1577
+ def test_series_add_sub_with_UInt64():
1578
+ # GH 22023
1579
+ series1 = Series([1, 2, 3])
1580
+ series2 = Series([2, 1, 3], dtype="UInt64")
1581
+
1582
+ result = series1 + series2
1583
+ expected = Series([3, 3, 6], dtype="Float64")
1584
+ tm.assert_series_equal(result, expected)
1585
+
1586
+ result = series1 - series2
1587
+ expected = Series([-1, 1, 0], dtype="Float64")
1588
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_object.py ADDED
@@ -0,0 +1,410 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arithmetic tests for DataFrame/Series/Index/Array classes that should
2
+ # behave identically.
3
+ # Specifically for object dtype
4
+ import datetime
5
+ from decimal import Decimal
6
+ import operator
7
+
8
+ import numpy as np
9
+ import pytest
10
+
11
+ import pandas.util._test_decorators as td
12
+
13
+ import pandas as pd
14
+ from pandas import (
15
+ Series,
16
+ Timestamp,
17
+ option_context,
18
+ )
19
+ import pandas._testing as tm
20
+ from pandas.core import ops
21
+
22
+ # ------------------------------------------------------------------
23
+ # Comparisons
24
+
25
+
26
+ class TestObjectComparisons:
27
+ def test_comparison_object_numeric_nas(self, comparison_op):
28
+ ser = Series(np.random.default_rng(2).standard_normal(10), dtype=object)
29
+ shifted = ser.shift(2)
30
+
31
+ func = comparison_op
32
+
33
+ result = func(ser, shifted)
34
+ expected = func(ser.astype(float), shifted.astype(float))
35
+ tm.assert_series_equal(result, expected)
36
+
37
+ @pytest.mark.parametrize(
38
+ "infer_string", [False, pytest.param(True, marks=td.skip_if_no("pyarrow"))]
39
+ )
40
+ def test_object_comparisons(self, infer_string):
41
+ with option_context("future.infer_string", infer_string):
42
+ ser = Series(["a", "b", np.nan, "c", "a"])
43
+
44
+ result = ser == "a"
45
+ expected = Series([True, False, False, False, True])
46
+ tm.assert_series_equal(result, expected)
47
+
48
+ result = ser < "a"
49
+ expected = Series([False, False, False, False, False])
50
+ tm.assert_series_equal(result, expected)
51
+
52
+ result = ser != "a"
53
+ expected = -(ser == "a")
54
+ tm.assert_series_equal(result, expected)
55
+
56
+ @pytest.mark.parametrize("dtype", [None, object])
57
+ def test_more_na_comparisons(self, dtype):
58
+ left = Series(["a", np.nan, "c"], dtype=dtype)
59
+ right = Series(["a", np.nan, "d"], dtype=dtype)
60
+
61
+ result = left == right
62
+ expected = Series([True, False, False])
63
+ tm.assert_series_equal(result, expected)
64
+
65
+ result = left != right
66
+ expected = Series([False, True, True])
67
+ tm.assert_series_equal(result, expected)
68
+
69
+ result = left == np.nan
70
+ expected = Series([False, False, False])
71
+ tm.assert_series_equal(result, expected)
72
+
73
+ result = left != np.nan
74
+ expected = Series([True, True, True])
75
+ tm.assert_series_equal(result, expected)
76
+
77
+
78
+ # ------------------------------------------------------------------
79
+ # Arithmetic
80
+
81
+
82
+ class TestArithmetic:
83
+ def test_add_period_to_array_of_offset(self):
84
+ # GH#50162
85
+ per = pd.Period("2012-1-1", freq="D")
86
+ pi = pd.period_range("2012-1-1", periods=10, freq="D")
87
+ idx = per - pi
88
+
89
+ expected = pd.Index([x + per for x in idx], dtype=object)
90
+ result = idx + per
91
+ tm.assert_index_equal(result, expected)
92
+
93
+ result = per + idx
94
+ tm.assert_index_equal(result, expected)
95
+
96
+ # TODO: parametrize
97
+ def test_pow_ops_object(self):
98
+ # GH#22922
99
+ # pow is weird with masking & 1, so testing here
100
+ a = Series([1, np.nan, 1, np.nan], dtype=object)
101
+ b = Series([1, np.nan, np.nan, 1], dtype=object)
102
+ result = a**b
103
+ expected = Series(a.values**b.values, dtype=object)
104
+ tm.assert_series_equal(result, expected)
105
+
106
+ result = b**a
107
+ expected = Series(b.values**a.values, dtype=object)
108
+
109
+ tm.assert_series_equal(result, expected)
110
+
111
+ @pytest.mark.parametrize("op", [operator.add, ops.radd])
112
+ @pytest.mark.parametrize("other", ["category", "Int64"])
113
+ def test_add_extension_scalar(self, other, box_with_array, op):
114
+ # GH#22378
115
+ # Check that scalars satisfying is_extension_array_dtype(obj)
116
+ # do not incorrectly try to dispatch to an ExtensionArray operation
117
+
118
+ arr = Series(["a", "b", "c"])
119
+ expected = Series([op(x, other) for x in arr])
120
+
121
+ arr = tm.box_expected(arr, box_with_array)
122
+ expected = tm.box_expected(expected, box_with_array)
123
+
124
+ result = op(arr, other)
125
+ tm.assert_equal(result, expected)
126
+
127
+ def test_objarr_add_str(self, box_with_array):
128
+ ser = Series(["x", np.nan, "x"])
129
+ expected = Series(["xa", np.nan, "xa"])
130
+
131
+ ser = tm.box_expected(ser, box_with_array)
132
+ expected = tm.box_expected(expected, box_with_array)
133
+
134
+ result = ser + "a"
135
+ tm.assert_equal(result, expected)
136
+
137
+ def test_objarr_radd_str(self, box_with_array):
138
+ ser = Series(["x", np.nan, "x"])
139
+ expected = Series(["ax", np.nan, "ax"])
140
+
141
+ ser = tm.box_expected(ser, box_with_array)
142
+ expected = tm.box_expected(expected, box_with_array)
143
+
144
+ result = "a" + ser
145
+ tm.assert_equal(result, expected)
146
+
147
+ @pytest.mark.parametrize(
148
+ "data",
149
+ [
150
+ [1, 2, 3],
151
+ [1.1, 2.2, 3.3],
152
+ [Timestamp("2011-01-01"), Timestamp("2011-01-02"), pd.NaT],
153
+ ["x", "y", 1],
154
+ ],
155
+ )
156
+ @pytest.mark.parametrize("dtype", [None, object])
157
+ def test_objarr_radd_str_invalid(self, dtype, data, box_with_array):
158
+ ser = Series(data, dtype=dtype)
159
+
160
+ ser = tm.box_expected(ser, box_with_array)
161
+ msg = "|".join(
162
+ [
163
+ "can only concatenate str",
164
+ "did not contain a loop with signature matching types",
165
+ "unsupported operand type",
166
+ "must be str",
167
+ ]
168
+ )
169
+ with pytest.raises(TypeError, match=msg):
170
+ "foo_" + ser
171
+
172
+ @pytest.mark.parametrize("op", [operator.add, ops.radd, operator.sub, ops.rsub])
173
+ def test_objarr_add_invalid(self, op, box_with_array):
174
+ # invalid ops
175
+ box = box_with_array
176
+
177
+ obj_ser = Series(list("abc"), dtype=object, name="objects")
178
+
179
+ obj_ser = tm.box_expected(obj_ser, box)
180
+ msg = "|".join(
181
+ [
182
+ "can only concatenate str",
183
+ "unsupported operand type",
184
+ "must be str",
185
+ "has no kernel",
186
+ ]
187
+ )
188
+ with pytest.raises(Exception, match=msg):
189
+ op(obj_ser, 1)
190
+ with pytest.raises(Exception, match=msg):
191
+ op(obj_ser, np.array(1, dtype=np.int64))
192
+
193
+ # TODO: Moved from tests.series.test_operators; needs cleanup
194
+ def test_operators_na_handling(self):
195
+ ser = Series(["foo", "bar", "baz", np.nan])
196
+ result = "prefix_" + ser
197
+ expected = Series(["prefix_foo", "prefix_bar", "prefix_baz", np.nan])
198
+ tm.assert_series_equal(result, expected)
199
+
200
+ result = ser + "_suffix"
201
+ expected = Series(["foo_suffix", "bar_suffix", "baz_suffix", np.nan])
202
+ tm.assert_series_equal(result, expected)
203
+
204
+ # TODO: parametrize over box
205
+ @pytest.mark.parametrize("dtype", [None, object])
206
+ def test_series_with_dtype_radd_timedelta(self, dtype):
207
+ # note this test is _not_ aimed at timedelta64-dtyped Series
208
+ # as of 2.0 we retain object dtype when ser.dtype == object
209
+ ser = Series(
210
+ [pd.Timedelta("1 days"), pd.Timedelta("2 days"), pd.Timedelta("3 days")],
211
+ dtype=dtype,
212
+ )
213
+ expected = Series(
214
+ [pd.Timedelta("4 days"), pd.Timedelta("5 days"), pd.Timedelta("6 days")],
215
+ dtype=dtype,
216
+ )
217
+
218
+ result = pd.Timedelta("3 days") + ser
219
+ tm.assert_series_equal(result, expected)
220
+
221
+ result = ser + pd.Timedelta("3 days")
222
+ tm.assert_series_equal(result, expected)
223
+
224
+ # TODO: cleanup & parametrize over box
225
+ def test_mixed_timezone_series_ops_object(self):
226
+ # GH#13043
227
+ ser = Series(
228
+ [
229
+ Timestamp("2015-01-01", tz="US/Eastern"),
230
+ Timestamp("2015-01-01", tz="Asia/Tokyo"),
231
+ ],
232
+ name="xxx",
233
+ )
234
+ assert ser.dtype == object
235
+
236
+ exp = Series(
237
+ [
238
+ Timestamp("2015-01-02", tz="US/Eastern"),
239
+ Timestamp("2015-01-02", tz="Asia/Tokyo"),
240
+ ],
241
+ name="xxx",
242
+ )
243
+ tm.assert_series_equal(ser + pd.Timedelta("1 days"), exp)
244
+ tm.assert_series_equal(pd.Timedelta("1 days") + ser, exp)
245
+
246
+ # object series & object series
247
+ ser2 = Series(
248
+ [
249
+ Timestamp("2015-01-03", tz="US/Eastern"),
250
+ Timestamp("2015-01-05", tz="Asia/Tokyo"),
251
+ ],
252
+ name="xxx",
253
+ )
254
+ assert ser2.dtype == object
255
+ exp = Series(
256
+ [pd.Timedelta("2 days"), pd.Timedelta("4 days")], name="xxx", dtype=object
257
+ )
258
+ tm.assert_series_equal(ser2 - ser, exp)
259
+ tm.assert_series_equal(ser - ser2, -exp)
260
+
261
+ ser = Series(
262
+ [pd.Timedelta("01:00:00"), pd.Timedelta("02:00:00")],
263
+ name="xxx",
264
+ dtype=object,
265
+ )
266
+ assert ser.dtype == object
267
+
268
+ exp = Series(
269
+ [pd.Timedelta("01:30:00"), pd.Timedelta("02:30:00")],
270
+ name="xxx",
271
+ dtype=object,
272
+ )
273
+ tm.assert_series_equal(ser + pd.Timedelta("00:30:00"), exp)
274
+ tm.assert_series_equal(pd.Timedelta("00:30:00") + ser, exp)
275
+
276
+ # TODO: cleanup & parametrize over box
277
+ def test_iadd_preserves_name(self):
278
+ # GH#17067, GH#19723 __iadd__ and __isub__ should preserve index name
279
+ ser = Series([1, 2, 3])
280
+ ser.index.name = "foo"
281
+
282
+ ser.index += 1
283
+ assert ser.index.name == "foo"
284
+
285
+ ser.index -= 1
286
+ assert ser.index.name == "foo"
287
+
288
+ def test_add_string(self):
289
+ # from bug report
290
+ index = pd.Index(["a", "b", "c"])
291
+ index2 = index + "foo"
292
+
293
+ assert "a" not in index2
294
+ assert "afoo" in index2
295
+
296
+ def test_iadd_string(self):
297
+ index = pd.Index(["a", "b", "c"])
298
+ # doesn't fail test unless there is a check before `+=`
299
+ assert "a" in index
300
+
301
+ index += "_x"
302
+ assert "a_x" in index
303
+
304
+ def test_add(self):
305
+ index = pd.Index([str(i) for i in range(10)])
306
+ expected = pd.Index(index.values * 2)
307
+ tm.assert_index_equal(index + index, expected)
308
+ tm.assert_index_equal(index + index.tolist(), expected)
309
+ tm.assert_index_equal(index.tolist() + index, expected)
310
+
311
+ # test add and radd
312
+ index = pd.Index(list("abc"))
313
+ expected = pd.Index(["a1", "b1", "c1"])
314
+ tm.assert_index_equal(index + "1", expected)
315
+ expected = pd.Index(["1a", "1b", "1c"])
316
+ tm.assert_index_equal("1" + index, expected)
317
+
318
+ def test_sub_fail(self):
319
+ index = pd.Index([str(i) for i in range(10)])
320
+
321
+ msg = "unsupported operand type|Cannot broadcast|sub' not supported"
322
+ with pytest.raises(TypeError, match=msg):
323
+ index - "a"
324
+ with pytest.raises(TypeError, match=msg):
325
+ index - index
326
+ with pytest.raises(TypeError, match=msg):
327
+ index - index.tolist()
328
+ with pytest.raises(TypeError, match=msg):
329
+ index.tolist() - index
330
+
331
+ def test_sub_object(self):
332
+ # GH#19369
333
+ index = pd.Index([Decimal(1), Decimal(2)])
334
+ expected = pd.Index([Decimal(0), Decimal(1)])
335
+
336
+ result = index - Decimal(1)
337
+ tm.assert_index_equal(result, expected)
338
+
339
+ result = index - pd.Index([Decimal(1), Decimal(1)])
340
+ tm.assert_index_equal(result, expected)
341
+
342
+ msg = "unsupported operand type"
343
+ with pytest.raises(TypeError, match=msg):
344
+ index - "foo"
345
+
346
+ with pytest.raises(TypeError, match=msg):
347
+ index - np.array([2, "foo"], dtype=object)
348
+
349
+ def test_rsub_object(self, fixed_now_ts):
350
+ # GH#19369
351
+ index = pd.Index([Decimal(1), Decimal(2)])
352
+ expected = pd.Index([Decimal(1), Decimal(0)])
353
+
354
+ result = Decimal(2) - index
355
+ tm.assert_index_equal(result, expected)
356
+
357
+ result = np.array([Decimal(2), Decimal(2)]) - index
358
+ tm.assert_index_equal(result, expected)
359
+
360
+ msg = "unsupported operand type"
361
+ with pytest.raises(TypeError, match=msg):
362
+ "foo" - index
363
+
364
+ with pytest.raises(TypeError, match=msg):
365
+ np.array([True, fixed_now_ts]) - index
366
+
367
+
368
+ class MyIndex(pd.Index):
369
+ # Simple index subclass that tracks ops calls.
370
+
371
+ _calls: int
372
+
373
+ @classmethod
374
+ def _simple_new(cls, values, name=None, dtype=None):
375
+ result = object.__new__(cls)
376
+ result._data = values
377
+ result._name = name
378
+ result._calls = 0
379
+ result._reset_identity()
380
+
381
+ return result
382
+
383
+ def __add__(self, other):
384
+ self._calls += 1
385
+ return self._simple_new(self._data)
386
+
387
+ def __radd__(self, other):
388
+ return self.__add__(other)
389
+
390
+
391
+ @pytest.mark.parametrize(
392
+ "other",
393
+ [
394
+ [datetime.timedelta(1), datetime.timedelta(2)],
395
+ [datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 2)],
396
+ [pd.Period("2000"), pd.Period("2001")],
397
+ ["a", "b"],
398
+ ],
399
+ ids=["timedelta", "datetime", "period", "object"],
400
+ )
401
+ def test_index_ops_defer_to_unknown_subclasses(other):
402
+ # https://github.com/pandas-dev/pandas/issues/31109
403
+ values = np.array(
404
+ [datetime.date(2000, 1, 1), datetime.date(2000, 1, 2)], dtype=object
405
+ )
406
+ a = MyIndex._simple_new(values)
407
+ other = pd.Index(other)
408
+ result = other + a
409
+ assert isinstance(result, MyIndex)
410
+ assert a._calls == 1
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_period.py ADDED
@@ -0,0 +1,1679 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arithmetic tests for DataFrame/Series/Index/Array classes that should
2
+ # behave identically.
3
+ # Specifically for Period dtype
4
+ import operator
5
+
6
+ import numpy as np
7
+ import pytest
8
+
9
+ from pandas._libs.tslibs import (
10
+ IncompatibleFrequency,
11
+ Period,
12
+ Timestamp,
13
+ to_offset,
14
+ )
15
+
16
+ import pandas as pd
17
+ from pandas import (
18
+ PeriodIndex,
19
+ Series,
20
+ Timedelta,
21
+ TimedeltaIndex,
22
+ period_range,
23
+ )
24
+ import pandas._testing as tm
25
+ from pandas.core import ops
26
+ from pandas.core.arrays import TimedeltaArray
27
+ from pandas.tests.arithmetic.common import (
28
+ assert_invalid_addsub_type,
29
+ assert_invalid_comparison,
30
+ get_upcast_box,
31
+ )
32
+
33
+ _common_mismatch = [
34
+ pd.offsets.YearBegin(2),
35
+ pd.offsets.MonthBegin(1),
36
+ pd.offsets.Minute(),
37
+ ]
38
+
39
+
40
+ @pytest.fixture(
41
+ params=[
42
+ Timedelta(minutes=30).to_pytimedelta(),
43
+ np.timedelta64(30, "s"),
44
+ Timedelta(seconds=30),
45
+ *_common_mismatch,
46
+ ]
47
+ )
48
+ def not_hourly(request):
49
+ """
50
+ Several timedelta-like and DateOffset instances that are _not_
51
+ compatible with Hourly frequencies.
52
+ """
53
+ return request.param
54
+
55
+
56
+ @pytest.fixture(
57
+ params=[
58
+ np.timedelta64(365, "D"),
59
+ Timedelta(days=365).to_pytimedelta(),
60
+ Timedelta(days=365),
61
+ *_common_mismatch,
62
+ ]
63
+ )
64
+ def mismatched_freq(request):
65
+ """
66
+ Several timedelta-like and DateOffset instances that are _not_
67
+ compatible with Monthly or Annual frequencies.
68
+ """
69
+ return request.param
70
+
71
+
72
+ # ------------------------------------------------------------------
73
+ # Comparisons
74
+
75
+
76
+ class TestPeriodArrayLikeComparisons:
77
+ # Comparison tests for PeriodDtype vectors fully parametrized over
78
+ # DataFrame/Series/PeriodIndex/PeriodArray. Ideally all comparison
79
+ # tests will eventually end up here.
80
+
81
+ @pytest.mark.parametrize("other", ["2017", Period("2017", freq="D")])
82
+ def test_eq_scalar(self, other, box_with_array):
83
+ idx = PeriodIndex(["2017", "2017", "2018"], freq="D")
84
+ idx = tm.box_expected(idx, box_with_array)
85
+ xbox = get_upcast_box(idx, other, True)
86
+
87
+ expected = np.array([True, True, False])
88
+ expected = tm.box_expected(expected, xbox)
89
+
90
+ result = idx == other
91
+
92
+ tm.assert_equal(result, expected)
93
+
94
+ def test_compare_zerodim(self, box_with_array):
95
+ # GH#26689 make sure we unbox zero-dimensional arrays
96
+
97
+ pi = period_range("2000", periods=4)
98
+ other = np.array(pi.to_numpy()[0])
99
+
100
+ pi = tm.box_expected(pi, box_with_array)
101
+ xbox = get_upcast_box(pi, other, True)
102
+
103
+ result = pi <= other
104
+ expected = np.array([True, False, False, False])
105
+ expected = tm.box_expected(expected, xbox)
106
+ tm.assert_equal(result, expected)
107
+
108
+ @pytest.mark.parametrize(
109
+ "scalar",
110
+ [
111
+ "foo",
112
+ Timestamp("2021-01-01"),
113
+ Timedelta(days=4),
114
+ 9,
115
+ 9.5,
116
+ 2000, # specifically don't consider 2000 to match Period("2000", "D")
117
+ False,
118
+ None,
119
+ ],
120
+ )
121
+ def test_compare_invalid_scalar(self, box_with_array, scalar):
122
+ # GH#28980
123
+ # comparison with scalar that cannot be interpreted as a Period
124
+ pi = period_range("2000", periods=4)
125
+ parr = tm.box_expected(pi, box_with_array)
126
+ assert_invalid_comparison(parr, scalar, box_with_array)
127
+
128
+ @pytest.mark.parametrize(
129
+ "other",
130
+ [
131
+ pd.date_range("2000", periods=4).array,
132
+ pd.timedelta_range("1D", periods=4).array,
133
+ np.arange(4),
134
+ np.arange(4).astype(np.float64),
135
+ list(range(4)),
136
+ # match Period semantics by not treating integers as Periods
137
+ [2000, 2001, 2002, 2003],
138
+ np.arange(2000, 2004),
139
+ np.arange(2000, 2004).astype(object),
140
+ pd.Index([2000, 2001, 2002, 2003]),
141
+ ],
142
+ )
143
+ def test_compare_invalid_listlike(self, box_with_array, other):
144
+ pi = period_range("2000", periods=4)
145
+ parr = tm.box_expected(pi, box_with_array)
146
+ assert_invalid_comparison(parr, other, box_with_array)
147
+
148
+ @pytest.mark.parametrize("other_box", [list, np.array, lambda x: x.astype(object)])
149
+ def test_compare_object_dtype(self, box_with_array, other_box):
150
+ pi = period_range("2000", periods=5)
151
+ parr = tm.box_expected(pi, box_with_array)
152
+
153
+ other = other_box(pi)
154
+ xbox = get_upcast_box(parr, other, True)
155
+
156
+ expected = np.array([True, True, True, True, True])
157
+ expected = tm.box_expected(expected, xbox)
158
+
159
+ result = parr == other
160
+ tm.assert_equal(result, expected)
161
+ result = parr <= other
162
+ tm.assert_equal(result, expected)
163
+ result = parr >= other
164
+ tm.assert_equal(result, expected)
165
+
166
+ result = parr != other
167
+ tm.assert_equal(result, ~expected)
168
+ result = parr < other
169
+ tm.assert_equal(result, ~expected)
170
+ result = parr > other
171
+ tm.assert_equal(result, ~expected)
172
+
173
+ other = other_box(pi[::-1])
174
+
175
+ expected = np.array([False, False, True, False, False])
176
+ expected = tm.box_expected(expected, xbox)
177
+ result = parr == other
178
+ tm.assert_equal(result, expected)
179
+
180
+ expected = np.array([True, True, True, False, False])
181
+ expected = tm.box_expected(expected, xbox)
182
+ result = parr <= other
183
+ tm.assert_equal(result, expected)
184
+
185
+ expected = np.array([False, False, True, True, True])
186
+ expected = tm.box_expected(expected, xbox)
187
+ result = parr >= other
188
+ tm.assert_equal(result, expected)
189
+
190
+ expected = np.array([True, True, False, True, True])
191
+ expected = tm.box_expected(expected, xbox)
192
+ result = parr != other
193
+ tm.assert_equal(result, expected)
194
+
195
+ expected = np.array([True, True, False, False, False])
196
+ expected = tm.box_expected(expected, xbox)
197
+ result = parr < other
198
+ tm.assert_equal(result, expected)
199
+
200
+ expected = np.array([False, False, False, True, True])
201
+ expected = tm.box_expected(expected, xbox)
202
+ result = parr > other
203
+ tm.assert_equal(result, expected)
204
+
205
+
206
+ class TestPeriodIndexComparisons:
207
+ # TODO: parameterize over boxes
208
+
209
+ def test_pi_cmp_period(self):
210
+ idx = period_range("2007-01", periods=20, freq="M")
211
+ per = idx[10]
212
+
213
+ result = idx < per
214
+ exp = idx.values < idx.values[10]
215
+ tm.assert_numpy_array_equal(result, exp)
216
+
217
+ # Tests Period.__richcmp__ against ndarray[object, ndim=2]
218
+ result = idx.values.reshape(10, 2) < per
219
+ tm.assert_numpy_array_equal(result, exp.reshape(10, 2))
220
+
221
+ # Tests Period.__richcmp__ against ndarray[object, ndim=0]
222
+ result = idx < np.array(per)
223
+ tm.assert_numpy_array_equal(result, exp)
224
+
225
+ # TODO: moved from test_datetime64; de-duplicate with version below
226
+ def test_parr_cmp_period_scalar2(self, box_with_array):
227
+ pi = period_range("2000-01-01", periods=10, freq="D")
228
+
229
+ val = pi[3]
230
+ expected = [x > val for x in pi]
231
+
232
+ ser = tm.box_expected(pi, box_with_array)
233
+ xbox = get_upcast_box(ser, val, True)
234
+
235
+ expected = tm.box_expected(expected, xbox)
236
+ result = ser > val
237
+ tm.assert_equal(result, expected)
238
+
239
+ val = pi[5]
240
+ result = ser > val
241
+ expected = [x > val for x in pi]
242
+ expected = tm.box_expected(expected, xbox)
243
+ tm.assert_equal(result, expected)
244
+
245
+ @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
246
+ def test_parr_cmp_period_scalar(self, freq, box_with_array):
247
+ # GH#13200
248
+ base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq)
249
+ base = tm.box_expected(base, box_with_array)
250
+ per = Period("2011-02", freq=freq)
251
+ xbox = get_upcast_box(base, per, True)
252
+
253
+ exp = np.array([False, True, False, False])
254
+ exp = tm.box_expected(exp, xbox)
255
+ tm.assert_equal(base == per, exp)
256
+ tm.assert_equal(per == base, exp)
257
+
258
+ exp = np.array([True, False, True, True])
259
+ exp = tm.box_expected(exp, xbox)
260
+ tm.assert_equal(base != per, exp)
261
+ tm.assert_equal(per != base, exp)
262
+
263
+ exp = np.array([False, False, True, True])
264
+ exp = tm.box_expected(exp, xbox)
265
+ tm.assert_equal(base > per, exp)
266
+ tm.assert_equal(per < base, exp)
267
+
268
+ exp = np.array([True, False, False, False])
269
+ exp = tm.box_expected(exp, xbox)
270
+ tm.assert_equal(base < per, exp)
271
+ tm.assert_equal(per > base, exp)
272
+
273
+ exp = np.array([False, True, True, True])
274
+ exp = tm.box_expected(exp, xbox)
275
+ tm.assert_equal(base >= per, exp)
276
+ tm.assert_equal(per <= base, exp)
277
+
278
+ exp = np.array([True, True, False, False])
279
+ exp = tm.box_expected(exp, xbox)
280
+ tm.assert_equal(base <= per, exp)
281
+ tm.assert_equal(per >= base, exp)
282
+
283
+ @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
284
+ def test_parr_cmp_pi(self, freq, box_with_array):
285
+ # GH#13200
286
+ base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq)
287
+ base = tm.box_expected(base, box_with_array)
288
+
289
+ # TODO: could also box idx?
290
+ idx = PeriodIndex(["2011-02", "2011-01", "2011-03", "2011-05"], freq=freq)
291
+
292
+ xbox = get_upcast_box(base, idx, True)
293
+
294
+ exp = np.array([False, False, True, False])
295
+ exp = tm.box_expected(exp, xbox)
296
+ tm.assert_equal(base == idx, exp)
297
+
298
+ exp = np.array([True, True, False, True])
299
+ exp = tm.box_expected(exp, xbox)
300
+ tm.assert_equal(base != idx, exp)
301
+
302
+ exp = np.array([False, True, False, False])
303
+ exp = tm.box_expected(exp, xbox)
304
+ tm.assert_equal(base > idx, exp)
305
+
306
+ exp = np.array([True, False, False, True])
307
+ exp = tm.box_expected(exp, xbox)
308
+ tm.assert_equal(base < idx, exp)
309
+
310
+ exp = np.array([False, True, True, False])
311
+ exp = tm.box_expected(exp, xbox)
312
+ tm.assert_equal(base >= idx, exp)
313
+
314
+ exp = np.array([True, False, True, True])
315
+ exp = tm.box_expected(exp, xbox)
316
+ tm.assert_equal(base <= idx, exp)
317
+
318
+ @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
319
+ def test_parr_cmp_pi_mismatched_freq(self, freq, box_with_array):
320
+ # GH#13200
321
+ # different base freq
322
+ base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq)
323
+ base = tm.box_expected(base, box_with_array)
324
+
325
+ msg = rf"Invalid comparison between dtype=period\[{freq}\] and Period"
326
+ with pytest.raises(TypeError, match=msg):
327
+ base <= Period("2011", freq="Y")
328
+
329
+ with pytest.raises(TypeError, match=msg):
330
+ Period("2011", freq="Y") >= base
331
+
332
+ # TODO: Could parametrize over boxes for idx?
333
+ idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="Y")
334
+ rev_msg = r"Invalid comparison between dtype=period\[Y-DEC\] and PeriodArray"
335
+ idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg
336
+ with pytest.raises(TypeError, match=idx_msg):
337
+ base <= idx
338
+
339
+ # Different frequency
340
+ msg = rf"Invalid comparison between dtype=period\[{freq}\] and Period"
341
+ with pytest.raises(TypeError, match=msg):
342
+ base <= Period("2011", freq="4M")
343
+
344
+ with pytest.raises(TypeError, match=msg):
345
+ Period("2011", freq="4M") >= base
346
+
347
+ idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="4M")
348
+ rev_msg = r"Invalid comparison between dtype=period\[4M\] and PeriodArray"
349
+ idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg
350
+ with pytest.raises(TypeError, match=idx_msg):
351
+ base <= idx
352
+
353
+ @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
354
+ def test_pi_cmp_nat(self, freq):
355
+ idx1 = PeriodIndex(["2011-01", "2011-02", "NaT", "2011-05"], freq=freq)
356
+ per = idx1[1]
357
+
358
+ result = idx1 > per
359
+ exp = np.array([False, False, False, True])
360
+ tm.assert_numpy_array_equal(result, exp)
361
+ result = per < idx1
362
+ tm.assert_numpy_array_equal(result, exp)
363
+
364
+ result = idx1 == pd.NaT
365
+ exp = np.array([False, False, False, False])
366
+ tm.assert_numpy_array_equal(result, exp)
367
+ result = pd.NaT == idx1
368
+ tm.assert_numpy_array_equal(result, exp)
369
+
370
+ result = idx1 != pd.NaT
371
+ exp = np.array([True, True, True, True])
372
+ tm.assert_numpy_array_equal(result, exp)
373
+ result = pd.NaT != idx1
374
+ tm.assert_numpy_array_equal(result, exp)
375
+
376
+ idx2 = PeriodIndex(["2011-02", "2011-01", "2011-04", "NaT"], freq=freq)
377
+ result = idx1 < idx2
378
+ exp = np.array([True, False, False, False])
379
+ tm.assert_numpy_array_equal(result, exp)
380
+
381
+ result = idx1 == idx2
382
+ exp = np.array([False, False, False, False])
383
+ tm.assert_numpy_array_equal(result, exp)
384
+
385
+ result = idx1 != idx2
386
+ exp = np.array([True, True, True, True])
387
+ tm.assert_numpy_array_equal(result, exp)
388
+
389
+ result = idx1 == idx1
390
+ exp = np.array([True, True, False, True])
391
+ tm.assert_numpy_array_equal(result, exp)
392
+
393
+ result = idx1 != idx1
394
+ exp = np.array([False, False, True, False])
395
+ tm.assert_numpy_array_equal(result, exp)
396
+
397
+ @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
398
+ def test_pi_cmp_nat_mismatched_freq_raises(self, freq):
399
+ idx1 = PeriodIndex(["2011-01", "2011-02", "NaT", "2011-05"], freq=freq)
400
+
401
+ diff = PeriodIndex(["2011-02", "2011-01", "2011-04", "NaT"], freq="4M")
402
+ msg = rf"Invalid comparison between dtype=period\[{freq}\] and PeriodArray"
403
+ with pytest.raises(TypeError, match=msg):
404
+ idx1 > diff
405
+
406
+ result = idx1 == diff
407
+ expected = np.array([False, False, False, False], dtype=bool)
408
+ tm.assert_numpy_array_equal(result, expected)
409
+
410
+ # TODO: De-duplicate with test_pi_cmp_nat
411
+ @pytest.mark.parametrize("dtype", [object, None])
412
+ def test_comp_nat(self, dtype):
413
+ left = PeriodIndex([Period("2011-01-01"), pd.NaT, Period("2011-01-03")])
414
+ right = PeriodIndex([pd.NaT, pd.NaT, Period("2011-01-03")])
415
+
416
+ if dtype is not None:
417
+ left = left.astype(dtype)
418
+ right = right.astype(dtype)
419
+
420
+ result = left == right
421
+ expected = np.array([False, False, True])
422
+ tm.assert_numpy_array_equal(result, expected)
423
+
424
+ result = left != right
425
+ expected = np.array([True, True, False])
426
+ tm.assert_numpy_array_equal(result, expected)
427
+
428
+ expected = np.array([False, False, False])
429
+ tm.assert_numpy_array_equal(left == pd.NaT, expected)
430
+ tm.assert_numpy_array_equal(pd.NaT == right, expected)
431
+
432
+ expected = np.array([True, True, True])
433
+ tm.assert_numpy_array_equal(left != pd.NaT, expected)
434
+ tm.assert_numpy_array_equal(pd.NaT != left, expected)
435
+
436
+ expected = np.array([False, False, False])
437
+ tm.assert_numpy_array_equal(left < pd.NaT, expected)
438
+ tm.assert_numpy_array_equal(pd.NaT > left, expected)
439
+
440
+
441
+ class TestPeriodSeriesComparisons:
442
+ def test_cmp_series_period_series_mixed_freq(self):
443
+ # GH#13200
444
+ base = Series(
445
+ [
446
+ Period("2011", freq="Y"),
447
+ Period("2011-02", freq="M"),
448
+ Period("2013", freq="Y"),
449
+ Period("2011-04", freq="M"),
450
+ ]
451
+ )
452
+
453
+ ser = Series(
454
+ [
455
+ Period("2012", freq="Y"),
456
+ Period("2011-01", freq="M"),
457
+ Period("2013", freq="Y"),
458
+ Period("2011-05", freq="M"),
459
+ ]
460
+ )
461
+
462
+ exp = Series([False, False, True, False])
463
+ tm.assert_series_equal(base == ser, exp)
464
+
465
+ exp = Series([True, True, False, True])
466
+ tm.assert_series_equal(base != ser, exp)
467
+
468
+ exp = Series([False, True, False, False])
469
+ tm.assert_series_equal(base > ser, exp)
470
+
471
+ exp = Series([True, False, False, True])
472
+ tm.assert_series_equal(base < ser, exp)
473
+
474
+ exp = Series([False, True, True, False])
475
+ tm.assert_series_equal(base >= ser, exp)
476
+
477
+ exp = Series([True, False, True, True])
478
+ tm.assert_series_equal(base <= ser, exp)
479
+
480
+
481
+ class TestPeriodIndexSeriesComparisonConsistency:
482
+ """Test PeriodIndex and Period Series Ops consistency"""
483
+
484
+ # TODO: needs parametrization+de-duplication
485
+
486
+ def _check(self, values, func, expected):
487
+ # Test PeriodIndex and Period Series Ops consistency
488
+
489
+ idx = PeriodIndex(values)
490
+ result = func(idx)
491
+
492
+ # check that we don't pass an unwanted type to tm.assert_equal
493
+ assert isinstance(expected, (pd.Index, np.ndarray))
494
+ tm.assert_equal(result, expected)
495
+
496
+ s = Series(values)
497
+ result = func(s)
498
+
499
+ exp = Series(expected, name=values.name)
500
+ tm.assert_series_equal(result, exp)
501
+
502
+ def test_pi_comp_period(self):
503
+ idx = PeriodIndex(
504
+ ["2011-01", "2011-02", "2011-03", "2011-04"], freq="M", name="idx"
505
+ )
506
+ per = idx[2]
507
+
508
+ f = lambda x: x == per
509
+ exp = np.array([False, False, True, False], dtype=np.bool_)
510
+ self._check(idx, f, exp)
511
+ f = lambda x: per == x
512
+ self._check(idx, f, exp)
513
+
514
+ f = lambda x: x != per
515
+ exp = np.array([True, True, False, True], dtype=np.bool_)
516
+ self._check(idx, f, exp)
517
+ f = lambda x: per != x
518
+ self._check(idx, f, exp)
519
+
520
+ f = lambda x: per >= x
521
+ exp = np.array([True, True, True, False], dtype=np.bool_)
522
+ self._check(idx, f, exp)
523
+
524
+ f = lambda x: x > per
525
+ exp = np.array([False, False, False, True], dtype=np.bool_)
526
+ self._check(idx, f, exp)
527
+
528
+ f = lambda x: per >= x
529
+ exp = np.array([True, True, True, False], dtype=np.bool_)
530
+ self._check(idx, f, exp)
531
+
532
+ def test_pi_comp_period_nat(self):
533
+ idx = PeriodIndex(
534
+ ["2011-01", "NaT", "2011-03", "2011-04"], freq="M", name="idx"
535
+ )
536
+ per = idx[2]
537
+
538
+ f = lambda x: x == per
539
+ exp = np.array([False, False, True, False], dtype=np.bool_)
540
+ self._check(idx, f, exp)
541
+ f = lambda x: per == x
542
+ self._check(idx, f, exp)
543
+
544
+ f = lambda x: x == pd.NaT
545
+ exp = np.array([False, False, False, False], dtype=np.bool_)
546
+ self._check(idx, f, exp)
547
+ f = lambda x: pd.NaT == x
548
+ self._check(idx, f, exp)
549
+
550
+ f = lambda x: x != per
551
+ exp = np.array([True, True, False, True], dtype=np.bool_)
552
+ self._check(idx, f, exp)
553
+ f = lambda x: per != x
554
+ self._check(idx, f, exp)
555
+
556
+ f = lambda x: x != pd.NaT
557
+ exp = np.array([True, True, True, True], dtype=np.bool_)
558
+ self._check(idx, f, exp)
559
+ f = lambda x: pd.NaT != x
560
+ self._check(idx, f, exp)
561
+
562
+ f = lambda x: per >= x
563
+ exp = np.array([True, False, True, False], dtype=np.bool_)
564
+ self._check(idx, f, exp)
565
+
566
+ f = lambda x: x < per
567
+ exp = np.array([True, False, False, False], dtype=np.bool_)
568
+ self._check(idx, f, exp)
569
+
570
+ f = lambda x: x > pd.NaT
571
+ exp = np.array([False, False, False, False], dtype=np.bool_)
572
+ self._check(idx, f, exp)
573
+
574
+ f = lambda x: pd.NaT >= x
575
+ exp = np.array([False, False, False, False], dtype=np.bool_)
576
+ self._check(idx, f, exp)
577
+
578
+
579
+ # ------------------------------------------------------------------
580
+ # Arithmetic
581
+
582
+
583
+ class TestPeriodFrameArithmetic:
584
+ def test_ops_frame_period(self):
585
+ # GH#13043
586
+ df = pd.DataFrame(
587
+ {
588
+ "A": [Period("2015-01", freq="M"), Period("2015-02", freq="M")],
589
+ "B": [Period("2014-01", freq="M"), Period("2014-02", freq="M")],
590
+ }
591
+ )
592
+ assert df["A"].dtype == "Period[M]"
593
+ assert df["B"].dtype == "Period[M]"
594
+
595
+ p = Period("2015-03", freq="M")
596
+ off = p.freq
597
+ # dtype will be object because of original dtype
598
+ exp = pd.DataFrame(
599
+ {
600
+ "A": np.array([2 * off, 1 * off], dtype=object),
601
+ "B": np.array([14 * off, 13 * off], dtype=object),
602
+ }
603
+ )
604
+ tm.assert_frame_equal(p - df, exp)
605
+ tm.assert_frame_equal(df - p, -1 * exp)
606
+
607
+ df2 = pd.DataFrame(
608
+ {
609
+ "A": [Period("2015-05", freq="M"), Period("2015-06", freq="M")],
610
+ "B": [Period("2015-05", freq="M"), Period("2015-06", freq="M")],
611
+ }
612
+ )
613
+ assert df2["A"].dtype == "Period[M]"
614
+ assert df2["B"].dtype == "Period[M]"
615
+
616
+ exp = pd.DataFrame(
617
+ {
618
+ "A": np.array([4 * off, 4 * off], dtype=object),
619
+ "B": np.array([16 * off, 16 * off], dtype=object),
620
+ }
621
+ )
622
+ tm.assert_frame_equal(df2 - df, exp)
623
+ tm.assert_frame_equal(df - df2, -1 * exp)
624
+
625
+
626
+ class TestPeriodIndexArithmetic:
627
+ # ---------------------------------------------------------------
628
+ # __add__/__sub__ with PeriodIndex
629
+ # PeriodIndex + other is defined for integers and timedelta-like others
630
+ # PeriodIndex - other is defined for integers, timedelta-like others,
631
+ # and PeriodIndex (with matching freq)
632
+
633
+ def test_parr_add_iadd_parr_raises(self, box_with_array):
634
+ rng = period_range("1/1/2000", freq="D", periods=5)
635
+ other = period_range("1/6/2000", freq="D", periods=5)
636
+ # TODO: parametrize over boxes for other?
637
+
638
+ rng = tm.box_expected(rng, box_with_array)
639
+ # An earlier implementation of PeriodIndex addition performed
640
+ # a set operation (union). This has since been changed to
641
+ # raise a TypeError. See GH#14164 and GH#13077 for historical
642
+ # reference.
643
+ msg = r"unsupported operand type\(s\) for \+: .* and .*"
644
+ with pytest.raises(TypeError, match=msg):
645
+ rng + other
646
+
647
+ with pytest.raises(TypeError, match=msg):
648
+ rng += other
649
+
650
+ def test_pi_sub_isub_pi(self):
651
+ # GH#20049
652
+ # For historical reference see GH#14164, GH#13077.
653
+ # PeriodIndex subtraction originally performed set difference,
654
+ # then changed to raise TypeError before being implemented in GH#20049
655
+ rng = period_range("1/1/2000", freq="D", periods=5)
656
+ other = period_range("1/6/2000", freq="D", periods=5)
657
+
658
+ off = rng.freq
659
+ expected = pd.Index([-5 * off] * 5)
660
+ result = rng - other
661
+ tm.assert_index_equal(result, expected)
662
+
663
+ rng -= other
664
+ tm.assert_index_equal(rng, expected)
665
+
666
+ def test_pi_sub_pi_with_nat(self):
667
+ rng = period_range("1/1/2000", freq="D", periods=5)
668
+ other = rng[1:].insert(0, pd.NaT)
669
+ assert other[1:].equals(rng[1:])
670
+
671
+ result = rng - other
672
+ off = rng.freq
673
+ expected = pd.Index([pd.NaT, 0 * off, 0 * off, 0 * off, 0 * off])
674
+ tm.assert_index_equal(result, expected)
675
+
676
+ def test_parr_sub_pi_mismatched_freq(self, box_with_array, box_with_array2):
677
+ rng = period_range("1/1/2000", freq="D", periods=5)
678
+ other = period_range("1/6/2000", freq="h", periods=5)
679
+
680
+ rng = tm.box_expected(rng, box_with_array)
681
+ other = tm.box_expected(other, box_with_array2)
682
+ msg = r"Input has different freq=[hD] from PeriodArray\(freq=[Dh]\)"
683
+ with pytest.raises(IncompatibleFrequency, match=msg):
684
+ rng - other
685
+
686
+ @pytest.mark.parametrize("n", [1, 2, 3, 4])
687
+ def test_sub_n_gt_1_ticks(self, tick_classes, n):
688
+ # GH 23878
689
+ p1_d = "19910905"
690
+ p2_d = "19920406"
691
+ p1 = PeriodIndex([p1_d], freq=tick_classes(n))
692
+ p2 = PeriodIndex([p2_d], freq=tick_classes(n))
693
+
694
+ expected = PeriodIndex([p2_d], freq=p2.freq.base) - PeriodIndex(
695
+ [p1_d], freq=p1.freq.base
696
+ )
697
+
698
+ tm.assert_index_equal((p2 - p1), expected)
699
+
700
+ @pytest.mark.parametrize("n", [1, 2, 3, 4])
701
+ @pytest.mark.parametrize(
702
+ "offset, kwd_name",
703
+ [
704
+ (pd.offsets.YearEnd, "month"),
705
+ (pd.offsets.QuarterEnd, "startingMonth"),
706
+ (pd.offsets.MonthEnd, None),
707
+ (pd.offsets.Week, "weekday"),
708
+ ],
709
+ )
710
+ def test_sub_n_gt_1_offsets(self, offset, kwd_name, n):
711
+ # GH 23878
712
+ kwds = {kwd_name: 3} if kwd_name is not None else {}
713
+ p1_d = "19910905"
714
+ p2_d = "19920406"
715
+ freq = offset(n, normalize=False, **kwds)
716
+ p1 = PeriodIndex([p1_d], freq=freq)
717
+ p2 = PeriodIndex([p2_d], freq=freq)
718
+
719
+ result = p2 - p1
720
+ expected = PeriodIndex([p2_d], freq=freq.base) - PeriodIndex(
721
+ [p1_d], freq=freq.base
722
+ )
723
+
724
+ tm.assert_index_equal(result, expected)
725
+
726
+ # -------------------------------------------------------------
727
+ # Invalid Operations
728
+
729
+ @pytest.mark.parametrize(
730
+ "other",
731
+ [
732
+ # datetime scalars
733
+ Timestamp("2016-01-01"),
734
+ Timestamp("2016-01-01").to_pydatetime(),
735
+ Timestamp("2016-01-01").to_datetime64(),
736
+ # datetime-like arrays
737
+ pd.date_range("2016-01-01", periods=3, freq="h"),
738
+ pd.date_range("2016-01-01", periods=3, tz="Europe/Brussels"),
739
+ pd.date_range("2016-01-01", periods=3, freq="s")._data,
740
+ pd.date_range("2016-01-01", periods=3, tz="Asia/Tokyo")._data,
741
+ # Miscellaneous invalid types
742
+ 3.14,
743
+ np.array([2.0, 3.0, 4.0]),
744
+ ],
745
+ )
746
+ def test_parr_add_sub_invalid(self, other, box_with_array):
747
+ # GH#23215
748
+ rng = period_range("1/1/2000", freq="D", periods=3)
749
+ rng = tm.box_expected(rng, box_with_array)
750
+
751
+ msg = "|".join(
752
+ [
753
+ r"(:?cannot add PeriodArray and .*)",
754
+ r"(:?cannot subtract .* from (:?a\s)?.*)",
755
+ r"(:?unsupported operand type\(s\) for \+: .* and .*)",
756
+ r"unsupported operand type\(s\) for [+-]: .* and .*",
757
+ ]
758
+ )
759
+ assert_invalid_addsub_type(rng, other, msg)
760
+ with pytest.raises(TypeError, match=msg):
761
+ rng + other
762
+ with pytest.raises(TypeError, match=msg):
763
+ other + rng
764
+ with pytest.raises(TypeError, match=msg):
765
+ rng - other
766
+ with pytest.raises(TypeError, match=msg):
767
+ other - rng
768
+
769
+ # -----------------------------------------------------------------
770
+ # __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]
771
+
772
+ def test_pi_add_sub_td64_array_non_tick_raises(self):
773
+ rng = period_range("1/1/2000", freq="Q", periods=3)
774
+ tdi = TimedeltaIndex(["-1 Day", "-1 Day", "-1 Day"])
775
+ tdarr = tdi.values
776
+
777
+ msg = r"Cannot add or subtract timedelta64\[ns\] dtype from period\[Q-DEC\]"
778
+ with pytest.raises(TypeError, match=msg):
779
+ rng + tdarr
780
+ with pytest.raises(TypeError, match=msg):
781
+ tdarr + rng
782
+
783
+ with pytest.raises(TypeError, match=msg):
784
+ rng - tdarr
785
+ msg = r"cannot subtract PeriodArray from TimedeltaArray"
786
+ with pytest.raises(TypeError, match=msg):
787
+ tdarr - rng
788
+
789
+ def test_pi_add_sub_td64_array_tick(self):
790
+ # PeriodIndex + Timedelta-like is allowed only with
791
+ # tick-like frequencies
792
+ rng = period_range("1/1/2000", freq="90D", periods=3)
793
+ tdi = TimedeltaIndex(["-1 Day", "-1 Day", "-1 Day"])
794
+ tdarr = tdi.values
795
+
796
+ expected = period_range("12/31/1999", freq="90D", periods=3)
797
+ result = rng + tdi
798
+ tm.assert_index_equal(result, expected)
799
+ result = rng + tdarr
800
+ tm.assert_index_equal(result, expected)
801
+ result = tdi + rng
802
+ tm.assert_index_equal(result, expected)
803
+ result = tdarr + rng
804
+ tm.assert_index_equal(result, expected)
805
+
806
+ expected = period_range("1/2/2000", freq="90D", periods=3)
807
+
808
+ result = rng - tdi
809
+ tm.assert_index_equal(result, expected)
810
+ result = rng - tdarr
811
+ tm.assert_index_equal(result, expected)
812
+
813
+ msg = r"cannot subtract .* from .*"
814
+ with pytest.raises(TypeError, match=msg):
815
+ tdarr - rng
816
+
817
+ with pytest.raises(TypeError, match=msg):
818
+ tdi - rng
819
+
820
+ @pytest.mark.parametrize("pi_freq", ["D", "W", "Q", "h"])
821
+ @pytest.mark.parametrize("tdi_freq", [None, "h"])
822
+ def test_parr_sub_td64array(self, box_with_array, tdi_freq, pi_freq):
823
+ box = box_with_array
824
+ xbox = box if box not in [pd.array, tm.to_array] else pd.Index
825
+
826
+ tdi = TimedeltaIndex(["1 hours", "2 hours"], freq=tdi_freq)
827
+ dti = Timestamp("2018-03-07 17:16:40") + tdi
828
+ pi = dti.to_period(pi_freq)
829
+
830
+ # TODO: parametrize over box for pi?
831
+ td64obj = tm.box_expected(tdi, box)
832
+
833
+ if pi_freq == "h":
834
+ result = pi - td64obj
835
+ expected = (pi.to_timestamp("s") - tdi).to_period(pi_freq)
836
+ expected = tm.box_expected(expected, xbox)
837
+ tm.assert_equal(result, expected)
838
+
839
+ # Subtract from scalar
840
+ result = pi[0] - td64obj
841
+ expected = (pi[0].to_timestamp("s") - tdi).to_period(pi_freq)
842
+ expected = tm.box_expected(expected, box)
843
+ tm.assert_equal(result, expected)
844
+
845
+ elif pi_freq == "D":
846
+ # Tick, but non-compatible
847
+ msg = (
848
+ "Cannot add/subtract timedelta-like from PeriodArray that is "
849
+ "not an integer multiple of the PeriodArray's freq."
850
+ )
851
+ with pytest.raises(IncompatibleFrequency, match=msg):
852
+ pi - td64obj
853
+
854
+ with pytest.raises(IncompatibleFrequency, match=msg):
855
+ pi[0] - td64obj
856
+
857
+ else:
858
+ # With non-Tick freq, we could not add timedelta64 array regardless
859
+ # of what its resolution is
860
+ msg = "Cannot add or subtract timedelta64"
861
+ with pytest.raises(TypeError, match=msg):
862
+ pi - td64obj
863
+ with pytest.raises(TypeError, match=msg):
864
+ pi[0] - td64obj
865
+
866
+ # -----------------------------------------------------------------
867
+ # operations with array/Index of DateOffset objects
868
+
869
+ @pytest.mark.parametrize("box", [np.array, pd.Index])
870
+ def test_pi_add_offset_array(self, performance_warning, box):
871
+ # GH#18849
872
+ pi = PeriodIndex([Period("2015Q1"), Period("2016Q2")])
873
+ offs = box(
874
+ [
875
+ pd.offsets.QuarterEnd(n=1, startingMonth=12),
876
+ pd.offsets.QuarterEnd(n=-2, startingMonth=12),
877
+ ]
878
+ )
879
+ expected = PeriodIndex([Period("2015Q2"), Period("2015Q4")]).astype(object)
880
+
881
+ with tm.assert_produces_warning(performance_warning):
882
+ res = pi + offs
883
+ tm.assert_index_equal(res, expected)
884
+
885
+ with tm.assert_produces_warning(performance_warning):
886
+ res2 = offs + pi
887
+ tm.assert_index_equal(res2, expected)
888
+
889
+ unanchored = np.array([pd.offsets.Hour(n=1), pd.offsets.Minute(n=-2)])
890
+ # addition/subtraction ops with incompatible offsets should issue
891
+ # a PerformanceWarning and _then_ raise a TypeError.
892
+ msg = r"Input cannot be converted to Period\(freq=Q-DEC\)"
893
+ with pytest.raises(IncompatibleFrequency, match=msg):
894
+ with tm.assert_produces_warning(performance_warning):
895
+ pi + unanchored
896
+ with pytest.raises(IncompatibleFrequency, match=msg):
897
+ with tm.assert_produces_warning(performance_warning):
898
+ unanchored + pi
899
+
900
+ @pytest.mark.parametrize("box", [np.array, pd.Index])
901
+ def test_pi_sub_offset_array(self, performance_warning, box):
902
+ # GH#18824
903
+ pi = PeriodIndex([Period("2015Q1"), Period("2016Q2")])
904
+ other = box(
905
+ [
906
+ pd.offsets.QuarterEnd(n=1, startingMonth=12),
907
+ pd.offsets.QuarterEnd(n=-2, startingMonth=12),
908
+ ]
909
+ )
910
+
911
+ expected = PeriodIndex([pi[n] - other[n] for n in range(len(pi))])
912
+ expected = expected.astype(object)
913
+
914
+ with tm.assert_produces_warning(performance_warning):
915
+ res = pi - other
916
+ tm.assert_index_equal(res, expected)
917
+
918
+ anchored = box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)])
919
+
920
+ # addition/subtraction ops with anchored offsets should issue
921
+ # a PerformanceWarning and _then_ raise a TypeError.
922
+ msg = r"Input has different freq=-1M from Period\(freq=Q-DEC\)"
923
+ with pytest.raises(IncompatibleFrequency, match=msg):
924
+ with tm.assert_produces_warning(performance_warning):
925
+ pi - anchored
926
+ with pytest.raises(IncompatibleFrequency, match=msg):
927
+ with tm.assert_produces_warning(performance_warning):
928
+ anchored - pi
929
+
930
+ def test_pi_add_iadd_int(self, one):
931
+ # Variants of `one` for #19012
932
+ rng = period_range("2000-01-01 09:00", freq="h", periods=10)
933
+ result = rng + one
934
+ expected = period_range("2000-01-01 10:00", freq="h", periods=10)
935
+ tm.assert_index_equal(result, expected)
936
+ rng += one
937
+ tm.assert_index_equal(rng, expected)
938
+
939
+ def test_pi_sub_isub_int(self, one):
940
+ """
941
+ PeriodIndex.__sub__ and __isub__ with several representations of
942
+ the integer 1, e.g. int, np.int64, np.uint8, ...
943
+ """
944
+ rng = period_range("2000-01-01 09:00", freq="h", periods=10)
945
+ result = rng - one
946
+ expected = period_range("2000-01-01 08:00", freq="h", periods=10)
947
+ tm.assert_index_equal(result, expected)
948
+ rng -= one
949
+ tm.assert_index_equal(rng, expected)
950
+
951
+ @pytest.mark.parametrize("five", [5, np.array(5, dtype=np.int64)])
952
+ def test_pi_sub_intlike(self, five):
953
+ rng = period_range("2007-01", periods=50)
954
+
955
+ result = rng - five
956
+ exp = rng + (-five)
957
+ tm.assert_index_equal(result, exp)
958
+
959
+ def test_pi_add_sub_int_array_freqn_gt1(self):
960
+ # GH#47209 test adding array of ints when freq.n > 1 matches
961
+ # scalar behavior
962
+ pi = period_range("2016-01-01", periods=10, freq="2D")
963
+ arr = np.arange(10)
964
+ result = pi + arr
965
+ expected = pd.Index([x + y for x, y in zip(pi, arr, strict=True)])
966
+ tm.assert_index_equal(result, expected)
967
+
968
+ result = pi - arr
969
+ expected = pd.Index([x - y for x, y in zip(pi, arr, strict=True)])
970
+ tm.assert_index_equal(result, expected)
971
+
972
+ def test_pi_sub_isub_offset(self):
973
+ # offset
974
+ # DateOffset
975
+ rng = period_range("2014", "2024", freq="Y")
976
+ result = rng - pd.offsets.YearEnd(5)
977
+ expected = period_range("2009", "2019", freq="Y")
978
+ tm.assert_index_equal(result, expected)
979
+ rng -= pd.offsets.YearEnd(5)
980
+ tm.assert_index_equal(rng, expected)
981
+
982
+ rng = period_range("2014-01", "2016-12", freq="M")
983
+ result = rng - pd.offsets.MonthEnd(5)
984
+ expected = period_range("2013-08", "2016-07", freq="M")
985
+ tm.assert_index_equal(result, expected)
986
+
987
+ rng -= pd.offsets.MonthEnd(5)
988
+ tm.assert_index_equal(rng, expected)
989
+
990
+ @pytest.mark.parametrize("transpose", [True, False])
991
+ def test_pi_add_offset_n_gt1(self, box_with_array, transpose):
992
+ # GH#23215
993
+ # add offset to PeriodIndex with freq.n > 1
994
+
995
+ per = Period("2016-01", freq="2M")
996
+ pi = PeriodIndex([per])
997
+
998
+ expected = PeriodIndex(["2016-03"], freq="2M")
999
+
1000
+ pi = tm.box_expected(pi, box_with_array, transpose=transpose)
1001
+ expected = tm.box_expected(expected, box_with_array, transpose=transpose)
1002
+
1003
+ result = pi + per.freq
1004
+ tm.assert_equal(result, expected)
1005
+
1006
+ result = per.freq + pi
1007
+ tm.assert_equal(result, expected)
1008
+
1009
+ def test_pi_add_offset_n_gt1_not_divisible(self, box_with_array):
1010
+ # GH#23215
1011
+ # PeriodIndex with freq.n > 1 add offset with offset.n % freq.n != 0
1012
+ pi = PeriodIndex(["2016-01"], freq="2M")
1013
+ expected = PeriodIndex(["2016-04"], freq="2M")
1014
+
1015
+ pi = tm.box_expected(pi, box_with_array)
1016
+ expected = tm.box_expected(expected, box_with_array)
1017
+
1018
+ result = pi + to_offset("3ME")
1019
+ tm.assert_equal(result, expected)
1020
+
1021
+ result = to_offset("3ME") + pi
1022
+ tm.assert_equal(result, expected)
1023
+
1024
+ # ---------------------------------------------------------------
1025
+ # __add__/__sub__ with integer arrays
1026
+
1027
+ @pytest.mark.parametrize("int_holder", [np.array, pd.Index])
1028
+ @pytest.mark.parametrize("op", [operator.add, ops.radd])
1029
+ def test_pi_add_intarray(self, int_holder, op):
1030
+ # GH#19959
1031
+ pi = PeriodIndex([Period("2015Q1"), Period("NaT")])
1032
+ other = int_holder([4, -1])
1033
+
1034
+ result = op(pi, other)
1035
+ expected = PeriodIndex([Period("2016Q1"), Period("NaT")])
1036
+ tm.assert_index_equal(result, expected)
1037
+
1038
+ @pytest.mark.parametrize("int_holder", [np.array, pd.Index])
1039
+ def test_pi_sub_intarray(self, int_holder):
1040
+ # GH#19959
1041
+ pi = PeriodIndex([Period("2015Q1"), Period("NaT")])
1042
+ other = int_holder([4, -1])
1043
+
1044
+ result = pi - other
1045
+ expected = PeriodIndex([Period("2014Q1"), Period("NaT")])
1046
+ tm.assert_index_equal(result, expected)
1047
+
1048
+ msg = r"bad operand type for unary -: 'PeriodArray'"
1049
+ with pytest.raises(TypeError, match=msg):
1050
+ other - pi
1051
+
1052
+ # ---------------------------------------------------------------
1053
+ # Timedelta-like (timedelta, timedelta64, Timedelta, Tick)
1054
+ # TODO: Some of these are misnomers because of non-Tick DateOffsets
1055
+
1056
+ def test_parr_add_timedeltalike_minute_gt1(self, three_days, box_with_array):
1057
+ # GH#23031 adding a time-delta-like offset to a PeriodArray that has
1058
+ # minute frequency with n != 1. A more general case is tested below
1059
+ # in test_pi_add_timedeltalike_tick_gt1, but here we write out the
1060
+ # expected result more explicitly.
1061
+ other = three_days
1062
+ rng = period_range("2014-05-01", periods=3, freq="2D")
1063
+ rng = tm.box_expected(rng, box_with_array)
1064
+
1065
+ expected = PeriodIndex(["2014-05-04", "2014-05-06", "2014-05-08"], freq="2D")
1066
+ expected = tm.box_expected(expected, box_with_array)
1067
+
1068
+ result = rng + other
1069
+ tm.assert_equal(result, expected)
1070
+
1071
+ result = other + rng
1072
+ tm.assert_equal(result, expected)
1073
+
1074
+ # subtraction
1075
+ expected = PeriodIndex(["2014-04-28", "2014-04-30", "2014-05-02"], freq="2D")
1076
+ expected = tm.box_expected(expected, box_with_array)
1077
+ result = rng - other
1078
+ tm.assert_equal(result, expected)
1079
+
1080
+ msg = "|".join(
1081
+ [
1082
+ r"bad operand type for unary -: 'PeriodArray'",
1083
+ r"cannot subtract PeriodArray from timedelta64\[[hD]\]",
1084
+ ]
1085
+ )
1086
+ with pytest.raises(TypeError, match=msg):
1087
+ other - rng
1088
+
1089
+ @pytest.mark.parametrize("freqstr", ["5ns", "5us", "5ms", "5s", "5min", "5h", "5D"])
1090
+ def test_parr_add_timedeltalike_tick_gt1(self, three_days, freqstr, box_with_array):
1091
+ # GH#23031 adding a time-delta-like offset to a PeriodArray that has
1092
+ # tick-like frequency with n != 1
1093
+ other = three_days
1094
+ rng = period_range("2014-05-01", periods=6, freq=freqstr)
1095
+ first = rng[0]
1096
+ rng = tm.box_expected(rng, box_with_array)
1097
+
1098
+ expected = period_range(first + other, periods=6, freq=freqstr)
1099
+ expected = tm.box_expected(expected, box_with_array)
1100
+
1101
+ result = rng + other
1102
+ tm.assert_equal(result, expected)
1103
+
1104
+ result = other + rng
1105
+ tm.assert_equal(result, expected)
1106
+
1107
+ # subtraction
1108
+ expected = period_range(first - other, periods=6, freq=freqstr)
1109
+ expected = tm.box_expected(expected, box_with_array)
1110
+ result = rng - other
1111
+ tm.assert_equal(result, expected)
1112
+ msg = "|".join(
1113
+ [
1114
+ r"bad operand type for unary -: 'PeriodArray'",
1115
+ r"cannot subtract PeriodArray from timedelta64\[[hD]\]",
1116
+ ]
1117
+ )
1118
+ with pytest.raises(TypeError, match=msg):
1119
+ other - rng
1120
+
1121
+ def test_pi_add_iadd_timedeltalike_daily(self, three_days):
1122
+ # Tick
1123
+ other = three_days
1124
+ rng = period_range("2014-05-01", "2014-05-15", freq="D")
1125
+ expected = period_range("2014-05-04", "2014-05-18", freq="D")
1126
+
1127
+ result = rng + other
1128
+ tm.assert_index_equal(result, expected)
1129
+
1130
+ rng += other
1131
+ tm.assert_index_equal(rng, expected)
1132
+
1133
+ def test_pi_sub_isub_timedeltalike_daily(self, three_days):
1134
+ # Tick-like 3 Days
1135
+ other = three_days
1136
+ rng = period_range("2014-05-01", "2014-05-15", freq="D")
1137
+ expected = period_range("2014-04-28", "2014-05-12", freq="D")
1138
+
1139
+ result = rng - other
1140
+ tm.assert_index_equal(result, expected)
1141
+
1142
+ rng -= other
1143
+ tm.assert_index_equal(rng, expected)
1144
+
1145
+ def test_parr_add_sub_timedeltalike_freq_mismatch_daily(
1146
+ self, not_daily, box_with_array
1147
+ ):
1148
+ other = not_daily
1149
+ rng = period_range("2014-05-01", "2014-05-15", freq="D")
1150
+ rng = tm.box_expected(rng, box_with_array)
1151
+
1152
+ msg = "|".join(
1153
+ [
1154
+ # non-timedelta-like DateOffset
1155
+ "Input has different freq(=.+)? from Period.*?\\(freq=D\\)",
1156
+ # timedelta/td64/Timedelta but not a multiple of 24H
1157
+ "Cannot add/subtract timedelta-like from PeriodArray that is "
1158
+ "not an integer multiple of the PeriodArray's freq.",
1159
+ ]
1160
+ )
1161
+ with pytest.raises(IncompatibleFrequency, match=msg):
1162
+ rng + other
1163
+ with pytest.raises(IncompatibleFrequency, match=msg):
1164
+ rng += other
1165
+ with pytest.raises(IncompatibleFrequency, match=msg):
1166
+ rng - other
1167
+ with pytest.raises(IncompatibleFrequency, match=msg):
1168
+ rng -= other
1169
+
1170
+ def test_pi_add_iadd_timedeltalike_hourly(self, two_hours):
1171
+ other = two_hours
1172
+ rng = period_range("2014-01-01 10:00", "2014-01-05 10:00", freq="h")
1173
+ expected = period_range("2014-01-01 12:00", "2014-01-05 12:00", freq="h")
1174
+
1175
+ result = rng + other
1176
+ tm.assert_index_equal(result, expected)
1177
+
1178
+ rng += other
1179
+ tm.assert_index_equal(rng, expected)
1180
+
1181
+ def test_parr_add_timedeltalike_mismatched_freq_hourly(
1182
+ self, not_hourly, box_with_array
1183
+ ):
1184
+ other = not_hourly
1185
+ rng = period_range("2014-01-01 10:00", "2014-01-05 10:00", freq="h")
1186
+ rng = tm.box_expected(rng, box_with_array)
1187
+ msg = "|".join(
1188
+ [
1189
+ # non-timedelta-like DateOffset
1190
+ "Input has different freq(=.+)? from Period.*?\\(freq=h\\)",
1191
+ # timedelta/td64/Timedelta but not a multiple of 24H
1192
+ "Cannot add/subtract timedelta-like from PeriodArray that is "
1193
+ "not an integer multiple of the PeriodArray's freq.",
1194
+ ]
1195
+ )
1196
+
1197
+ with pytest.raises(IncompatibleFrequency, match=msg):
1198
+ rng + other
1199
+
1200
+ with pytest.raises(IncompatibleFrequency, match=msg):
1201
+ rng += other
1202
+
1203
+ def test_pi_sub_isub_timedeltalike_hourly(self, two_hours):
1204
+ other = two_hours
1205
+ rng = period_range("2014-01-01 10:00", "2014-01-05 10:00", freq="h")
1206
+ expected = period_range("2014-01-01 08:00", "2014-01-05 08:00", freq="h")
1207
+
1208
+ result = rng - other
1209
+ tm.assert_index_equal(result, expected)
1210
+
1211
+ rng -= other
1212
+ tm.assert_index_equal(rng, expected)
1213
+
1214
+ def test_add_iadd_timedeltalike_annual(self):
1215
+ # offset
1216
+ # DateOffset
1217
+ rng = period_range("2014", "2024", freq="Y")
1218
+ result = rng + pd.offsets.YearEnd(5)
1219
+ expected = period_range("2019", "2029", freq="Y")
1220
+ tm.assert_index_equal(result, expected)
1221
+ rng += pd.offsets.YearEnd(5)
1222
+ tm.assert_index_equal(rng, expected)
1223
+
1224
+ def test_pi_add_sub_timedeltalike_freq_mismatch_annual(self, mismatched_freq):
1225
+ other = mismatched_freq
1226
+ rng = period_range("2014", "2024", freq="Y")
1227
+ msg = "Input has different freq(=.+)? from Period.*?\\(freq=Y-DEC\\)"
1228
+ with pytest.raises(IncompatibleFrequency, match=msg):
1229
+ rng + other
1230
+ with pytest.raises(IncompatibleFrequency, match=msg):
1231
+ rng += other
1232
+ with pytest.raises(IncompatibleFrequency, match=msg):
1233
+ rng - other
1234
+ with pytest.raises(IncompatibleFrequency, match=msg):
1235
+ rng -= other
1236
+
1237
+ def test_pi_add_iadd_timedeltalike_M(self):
1238
+ rng = period_range("2014-01", "2016-12", freq="M")
1239
+ expected = period_range("2014-06", "2017-05", freq="M")
1240
+
1241
+ result = rng + pd.offsets.MonthEnd(5)
1242
+ tm.assert_index_equal(result, expected)
1243
+
1244
+ rng += pd.offsets.MonthEnd(5)
1245
+ tm.assert_index_equal(rng, expected)
1246
+
1247
+ def test_pi_add_sub_timedeltalike_freq_mismatch_monthly(self, mismatched_freq):
1248
+ other = mismatched_freq
1249
+ rng = period_range("2014-01", "2016-12", freq="M")
1250
+ msg = "Input has different freq(=.+)? from Period.*?\\(freq=M\\)"
1251
+ with pytest.raises(IncompatibleFrequency, match=msg):
1252
+ rng + other
1253
+ with pytest.raises(IncompatibleFrequency, match=msg):
1254
+ rng += other
1255
+ with pytest.raises(IncompatibleFrequency, match=msg):
1256
+ rng - other
1257
+ with pytest.raises(IncompatibleFrequency, match=msg):
1258
+ rng -= other
1259
+
1260
+ @pytest.mark.parametrize("transpose", [True, False])
1261
+ def test_parr_add_sub_td64_nat(self, box_with_array, transpose):
1262
+ # GH#23320 special handling for timedelta64("NaT")
1263
+ pi = period_range("1994-04-01", periods=9, freq="19D")
1264
+ other = np.timedelta64("NaT", "ns")
1265
+ expected = PeriodIndex(["NaT"] * 9, freq="19D")
1266
+
1267
+ obj = tm.box_expected(pi, box_with_array, transpose=transpose)
1268
+ expected = tm.box_expected(expected, box_with_array, transpose=transpose)
1269
+
1270
+ result = obj + other
1271
+ tm.assert_equal(result, expected)
1272
+ result = other + obj
1273
+ tm.assert_equal(result, expected)
1274
+ result = obj - other
1275
+ tm.assert_equal(result, expected)
1276
+ msg = r"cannot subtract .* from .*"
1277
+ with pytest.raises(TypeError, match=msg):
1278
+ other - obj
1279
+
1280
+ @pytest.mark.parametrize(
1281
+ "other",
1282
+ [
1283
+ np.array(["NaT"] * 9, dtype="m8[ns]"),
1284
+ TimedeltaArray._from_sequence(["NaT"] * 9, dtype="m8[ns]"),
1285
+ ],
1286
+ )
1287
+ def test_parr_add_sub_tdt64_nat_array(self, box_with_array, other):
1288
+ pi = period_range("1994-04-01", periods=9, freq="19D")
1289
+ expected = PeriodIndex(["NaT"] * 9, freq="19D")
1290
+
1291
+ obj = tm.box_expected(pi, box_with_array)
1292
+ expected = tm.box_expected(expected, box_with_array)
1293
+
1294
+ result = obj + other
1295
+ tm.assert_equal(result, expected)
1296
+ result = other + obj
1297
+ tm.assert_equal(result, expected)
1298
+ result = obj - other
1299
+ tm.assert_equal(result, expected)
1300
+ msg = r"cannot subtract .* from .*"
1301
+ with pytest.raises(TypeError, match=msg):
1302
+ other - obj
1303
+
1304
+ # some but not *all* NaT
1305
+ other = other.copy()
1306
+ other[0] = np.timedelta64(0, "ns")
1307
+ expected = PeriodIndex([pi[0]] + ["NaT"] * 8, freq="19D")
1308
+ expected = tm.box_expected(expected, box_with_array)
1309
+
1310
+ result = obj + other
1311
+ tm.assert_equal(result, expected)
1312
+ result = other + obj
1313
+ tm.assert_equal(result, expected)
1314
+ result = obj - other
1315
+ tm.assert_equal(result, expected)
1316
+ with pytest.raises(TypeError, match=msg):
1317
+ other - obj
1318
+
1319
+ # ---------------------------------------------------------------
1320
+ # Unsorted
1321
+
1322
+ def test_parr_add_sub_index(self):
1323
+ # Check that PeriodArray defers to Index on arithmetic ops
1324
+ pi = period_range("2000-12-31", periods=3)
1325
+ parr = pi.array
1326
+
1327
+ result = parr - pi
1328
+ expected = pi - pi
1329
+ tm.assert_index_equal(result, expected)
1330
+
1331
+ def test_parr_add_sub_object_array(self, performance_warning):
1332
+ pi = period_range("2000-12-31", periods=3, freq="D")
1333
+ parr = pi.array
1334
+
1335
+ other = np.array([Timedelta(days=1), pd.offsets.Day(2), 3])
1336
+
1337
+ with tm.assert_produces_warning(performance_warning):
1338
+ result = parr + other
1339
+
1340
+ expected = PeriodIndex(
1341
+ ["2001-01-01", "2001-01-03", "2001-01-05"], freq="D"
1342
+ )._data.astype(object)
1343
+ tm.assert_equal(result, expected)
1344
+
1345
+ with tm.assert_produces_warning(performance_warning):
1346
+ result = parr - other
1347
+
1348
+ expected = PeriodIndex(["2000-12-30"] * 3, freq="D")._data.astype(object)
1349
+ tm.assert_equal(result, expected)
1350
+
1351
+ def test_period_add_timestamp_raises(self, box_with_array):
1352
+ # GH#17983
1353
+ ts = Timestamp("2017")
1354
+ per = Period("2017", freq="M")
1355
+
1356
+ arr = pd.Index([per], dtype="Period[M]")
1357
+ arr = tm.box_expected(arr, box_with_array)
1358
+
1359
+ msg = "cannot add PeriodArray and Timestamp"
1360
+ with pytest.raises(TypeError, match=msg):
1361
+ arr + ts
1362
+ with pytest.raises(TypeError, match=msg):
1363
+ ts + arr
1364
+
1365
+ msg = "cannot add PeriodArray and DatetimeArray"
1366
+ with pytest.raises(TypeError, match=msg):
1367
+ arr + Series([ts])
1368
+ with pytest.raises(TypeError, match=msg):
1369
+ Series([ts]) + arr
1370
+ with pytest.raises(TypeError, match=msg):
1371
+ arr + pd.Index([ts])
1372
+ with pytest.raises(TypeError, match=msg):
1373
+ pd.Index([ts]) + arr
1374
+
1375
+ if box_with_array is pd.DataFrame:
1376
+ msg = "cannot add PeriodArray and DatetimeArray"
1377
+ else:
1378
+ msg = r"unsupported operand type\(s\) for \+: 'Period' and 'DatetimeArray"
1379
+ with pytest.raises(TypeError, match=msg):
1380
+ arr + pd.DataFrame([ts])
1381
+ if box_with_array is pd.DataFrame:
1382
+ msg = "cannot add PeriodArray and DatetimeArray"
1383
+ else:
1384
+ msg = r"unsupported operand type\(s\) for \+: 'DatetimeArray' and 'Period'"
1385
+ with pytest.raises(TypeError, match=msg):
1386
+ pd.DataFrame([ts]) + arr
1387
+
1388
+
1389
+ class TestPeriodSeriesArithmetic:
1390
+ def test_parr_add_timedeltalike_scalar(self, three_days, box_with_array):
1391
+ # GH#13043
1392
+ ser = Series(
1393
+ [Period("2015-01-01", freq="D"), Period("2015-01-02", freq="D")],
1394
+ name="xxx",
1395
+ )
1396
+ assert ser.dtype == "Period[D]"
1397
+
1398
+ expected = Series(
1399
+ [Period("2015-01-04", freq="D"), Period("2015-01-05", freq="D")],
1400
+ name="xxx",
1401
+ )
1402
+
1403
+ obj = tm.box_expected(ser, box_with_array)
1404
+ if box_with_array is pd.DataFrame:
1405
+ assert (obj.dtypes == "Period[D]").all()
1406
+
1407
+ expected = tm.box_expected(expected, box_with_array)
1408
+
1409
+ result = obj + three_days
1410
+ tm.assert_equal(result, expected)
1411
+
1412
+ result = three_days + obj
1413
+ tm.assert_equal(result, expected)
1414
+
1415
+ def test_ops_series_period(self):
1416
+ # GH#13043
1417
+ ser = Series(
1418
+ [Period("2015-01-01", freq="D"), Period("2015-01-02", freq="D")],
1419
+ name="xxx",
1420
+ )
1421
+ assert ser.dtype == "Period[D]"
1422
+
1423
+ per = Period("2015-01-10", freq="D")
1424
+ off = per.freq
1425
+ # dtype will be object because of original dtype
1426
+ expected = Series([9 * off, 8 * off], name="xxx", dtype=object)
1427
+ tm.assert_series_equal(per - ser, expected)
1428
+ tm.assert_series_equal(ser - per, -1 * expected)
1429
+
1430
+ s2 = Series(
1431
+ [Period("2015-01-05", freq="D"), Period("2015-01-04", freq="D")],
1432
+ name="xxx",
1433
+ )
1434
+ assert s2.dtype == "Period[D]"
1435
+
1436
+ expected = Series([4 * off, 2 * off], name="xxx", dtype=object)
1437
+ tm.assert_series_equal(s2 - ser, expected)
1438
+ tm.assert_series_equal(ser - s2, -1 * expected)
1439
+
1440
+
1441
+ class TestPeriodIndexSeriesMethods:
1442
+ """Test PeriodIndex and Period Series Ops consistency"""
1443
+
1444
+ def _check(self, values, func, expected):
1445
+ idx = PeriodIndex(values)
1446
+ result = func(idx)
1447
+ tm.assert_equal(result, expected)
1448
+
1449
+ ser = Series(values)
1450
+ result = func(ser)
1451
+
1452
+ exp = Series(expected, name=values.name)
1453
+ tm.assert_series_equal(result, exp)
1454
+
1455
+ def test_pi_ops(self):
1456
+ idx = PeriodIndex(
1457
+ ["2011-01", "2011-02", "2011-03", "2011-04"], freq="M", name="idx"
1458
+ )
1459
+
1460
+ expected = PeriodIndex(
1461
+ ["2011-03", "2011-04", "2011-05", "2011-06"], freq="M", name="idx"
1462
+ )
1463
+
1464
+ self._check(idx, lambda x: x + 2, expected)
1465
+ self._check(idx, lambda x: 2 + x, expected)
1466
+
1467
+ self._check(idx + 2, lambda x: x - 2, idx)
1468
+
1469
+ result = idx - Period("2011-01", freq="M")
1470
+ off = idx.freq
1471
+ exp = pd.Index([0 * off, 1 * off, 2 * off, 3 * off], name="idx")
1472
+ tm.assert_index_equal(result, exp)
1473
+
1474
+ result = Period("2011-01", freq="M") - idx
1475
+ exp = pd.Index([0 * off, -1 * off, -2 * off, -3 * off], name="idx")
1476
+ tm.assert_index_equal(result, exp)
1477
+
1478
+ @pytest.mark.parametrize("ng", ["str", 1.5])
1479
+ @pytest.mark.parametrize(
1480
+ "func",
1481
+ [
1482
+ lambda obj, ng: obj + ng,
1483
+ lambda obj, ng: ng + obj,
1484
+ lambda obj, ng: obj - ng,
1485
+ lambda obj, ng: ng - obj,
1486
+ lambda obj, ng: np.add(obj, ng),
1487
+ lambda obj, ng: np.add(ng, obj),
1488
+ lambda obj, ng: np.subtract(obj, ng),
1489
+ lambda obj, ng: np.subtract(ng, obj),
1490
+ ],
1491
+ )
1492
+ def test_parr_ops_errors(self, ng, func, box_with_array):
1493
+ idx = PeriodIndex(
1494
+ ["2011-01", "2011-02", "2011-03", "2011-04"], freq="M", name="idx"
1495
+ )
1496
+ obj = tm.box_expected(idx, box_with_array)
1497
+ msg = "|".join(
1498
+ [
1499
+ r"unsupported operand type\(s\)",
1500
+ "can only concatenate",
1501
+ r"must be str",
1502
+ "object to str implicitly",
1503
+ ]
1504
+ )
1505
+
1506
+ with pytest.raises(TypeError, match=msg):
1507
+ func(obj, ng)
1508
+
1509
+ def test_pi_ops_nat(self):
1510
+ idx = PeriodIndex(
1511
+ ["2011-01", "2011-02", "NaT", "2011-04"], freq="M", name="idx"
1512
+ )
1513
+ expected = PeriodIndex(
1514
+ ["2011-03", "2011-04", "NaT", "2011-06"], freq="M", name="idx"
1515
+ )
1516
+
1517
+ self._check(idx, lambda x: x + 2, expected)
1518
+ self._check(idx, lambda x: 2 + x, expected)
1519
+ self._check(idx, lambda x: np.add(x, 2), expected)
1520
+
1521
+ self._check(idx + 2, lambda x: x - 2, idx)
1522
+ self._check(idx + 2, lambda x: np.subtract(x, 2), idx)
1523
+
1524
+ # freq with mult
1525
+ idx = PeriodIndex(
1526
+ ["2011-01", "2011-02", "NaT", "2011-04"], freq="2M", name="idx"
1527
+ )
1528
+ expected = PeriodIndex(
1529
+ ["2011-07", "2011-08", "NaT", "2011-10"], freq="2M", name="idx"
1530
+ )
1531
+
1532
+ self._check(idx, lambda x: x + 3, expected)
1533
+ self._check(idx, lambda x: 3 + x, expected)
1534
+ self._check(idx, lambda x: np.add(x, 3), expected)
1535
+
1536
+ self._check(idx + 3, lambda x: x - 3, idx)
1537
+ self._check(idx + 3, lambda x: np.subtract(x, 3), idx)
1538
+
1539
+ def test_pi_ops_array_int(self):
1540
+ idx = PeriodIndex(
1541
+ ["2011-01", "2011-02", "NaT", "2011-04"], freq="M", name="idx"
1542
+ )
1543
+ f = lambda x: x + np.array([1, 2, 3, 4])
1544
+ exp = PeriodIndex(
1545
+ ["2011-02", "2011-04", "NaT", "2011-08"], freq="M", name="idx"
1546
+ )
1547
+ self._check(idx, f, exp)
1548
+
1549
+ f = lambda x: np.add(x, np.array([4, -1, 1, 2]))
1550
+ exp = PeriodIndex(
1551
+ ["2011-05", "2011-01", "NaT", "2011-06"], freq="M", name="idx"
1552
+ )
1553
+ self._check(idx, f, exp)
1554
+
1555
+ f = lambda x: x - np.array([1, 2, 3, 4])
1556
+ exp = PeriodIndex(
1557
+ ["2010-12", "2010-12", "NaT", "2010-12"], freq="M", name="idx"
1558
+ )
1559
+ self._check(idx, f, exp)
1560
+
1561
+ f = lambda x: np.subtract(x, np.array([3, 2, 3, -2]))
1562
+ exp = PeriodIndex(
1563
+ ["2010-10", "2010-12", "NaT", "2011-06"], freq="M", name="idx"
1564
+ )
1565
+ self._check(idx, f, exp)
1566
+
1567
+ def test_pi_ops_offset(self):
1568
+ idx = PeriodIndex(
1569
+ ["2011-01-01", "2011-02-01", "2011-03-01", "2011-04-01"],
1570
+ freq="D",
1571
+ name="idx",
1572
+ )
1573
+ f = lambda x: x + pd.offsets.Day()
1574
+ exp = PeriodIndex(
1575
+ ["2011-01-02", "2011-02-02", "2011-03-02", "2011-04-02"],
1576
+ freq="D",
1577
+ name="idx",
1578
+ )
1579
+ self._check(idx, f, exp)
1580
+
1581
+ f = lambda x: x + pd.offsets.Day(2)
1582
+ exp = PeriodIndex(
1583
+ ["2011-01-03", "2011-02-03", "2011-03-03", "2011-04-03"],
1584
+ freq="D",
1585
+ name="idx",
1586
+ )
1587
+ self._check(idx, f, exp)
1588
+
1589
+ f = lambda x: x - pd.offsets.Day(2)
1590
+ exp = PeriodIndex(
1591
+ ["2010-12-30", "2011-01-30", "2011-02-27", "2011-03-30"],
1592
+ freq="D",
1593
+ name="idx",
1594
+ )
1595
+ self._check(idx, f, exp)
1596
+
1597
+ def test_pi_offset_errors(self):
1598
+ idx = PeriodIndex(
1599
+ ["2011-01-01", "2011-02-01", "2011-03-01", "2011-04-01"],
1600
+ freq="D",
1601
+ name="idx",
1602
+ )
1603
+ ser = Series(idx)
1604
+
1605
+ msg = (
1606
+ "Cannot add/subtract timedelta-like from PeriodArray that is not "
1607
+ "an integer multiple of the PeriodArray's freq"
1608
+ )
1609
+ for obj in [idx, ser]:
1610
+ with pytest.raises(IncompatibleFrequency, match=msg):
1611
+ obj + pd.offsets.Hour(2)
1612
+
1613
+ with pytest.raises(IncompatibleFrequency, match=msg):
1614
+ pd.offsets.Hour(2) + obj
1615
+
1616
+ with pytest.raises(IncompatibleFrequency, match=msg):
1617
+ obj - pd.offsets.Hour(2)
1618
+
1619
+ def test_pi_sub_period(self):
1620
+ # GH#13071
1621
+ idx = PeriodIndex(
1622
+ ["2011-01", "2011-02", "2011-03", "2011-04"], freq="M", name="idx"
1623
+ )
1624
+
1625
+ result = idx - Period("2012-01", freq="M")
1626
+ off = idx.freq
1627
+ exp = pd.Index([-12 * off, -11 * off, -10 * off, -9 * off], name="idx")
1628
+ tm.assert_index_equal(result, exp)
1629
+
1630
+ result = np.subtract(idx, Period("2012-01", freq="M"))
1631
+ tm.assert_index_equal(result, exp)
1632
+
1633
+ result = Period("2012-01", freq="M") - idx
1634
+ exp = pd.Index([12 * off, 11 * off, 10 * off, 9 * off], name="idx")
1635
+ tm.assert_index_equal(result, exp)
1636
+
1637
+ result = np.subtract(Period("2012-01", freq="M"), idx)
1638
+ tm.assert_index_equal(result, exp)
1639
+
1640
+ exp = TimedeltaIndex(
1641
+ [np.nan, np.nan, np.nan, np.nan], name="idx", dtype="m8[ns]"
1642
+ )
1643
+ result = idx - Period("NaT", freq="M")
1644
+ tm.assert_index_equal(result, exp)
1645
+ assert result.freq == exp.freq
1646
+
1647
+ result = Period("NaT", freq="M") - idx
1648
+ tm.assert_index_equal(result, exp)
1649
+ assert result.freq == exp.freq
1650
+
1651
+ def test_pi_sub_pdnat(self):
1652
+ # GH#13071, GH#19389
1653
+ idx = PeriodIndex(
1654
+ ["2011-01", "2011-02", "NaT", "2011-04"], freq="M", name="idx"
1655
+ )
1656
+ exp = TimedeltaIndex([pd.NaT] * 4, name="idx", dtype="m8[ns]")
1657
+ tm.assert_index_equal(pd.NaT - idx, exp)
1658
+ tm.assert_index_equal(idx - pd.NaT, exp)
1659
+
1660
+ def test_pi_sub_period_nat(self):
1661
+ # GH#13071
1662
+ idx = PeriodIndex(
1663
+ ["2011-01", "NaT", "2011-03", "2011-04"], freq="M", name="idx"
1664
+ )
1665
+
1666
+ result = idx - Period("2012-01", freq="M")
1667
+ off = idx.freq
1668
+ exp = pd.Index([-12 * off, pd.NaT, -10 * off, -9 * off], name="idx")
1669
+ tm.assert_index_equal(result, exp)
1670
+
1671
+ result = Period("2012-01", freq="M") - idx
1672
+ exp = pd.Index([12 * off, pd.NaT, 10 * off, 9 * off], name="idx")
1673
+ tm.assert_index_equal(result, exp)
1674
+
1675
+ exp = TimedeltaIndex(
1676
+ [np.nan, np.nan, np.nan, np.nan], name="idx", dtype="m8[ns]"
1677
+ )
1678
+ tm.assert_index_equal(idx - Period("NaT", freq="M"), exp)
1679
+ tm.assert_index_equal(Period("NaT", freq="M") - idx, exp)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_string.py ADDED
@@ -0,0 +1,526 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import operator
2
+ from pathlib import Path
3
+
4
+ import numpy as np
5
+ import pytest
6
+
7
+ from pandas.compat import HAS_PYARROW
8
+ from pandas.errors import Pandas4Warning
9
+ import pandas.util._test_decorators as td
10
+
11
+ import pandas as pd
12
+ from pandas import (
13
+ NA,
14
+ ArrowDtype,
15
+ Series,
16
+ StringDtype,
17
+ )
18
+ import pandas._testing as tm
19
+ from pandas.core.construction import extract_array
20
+
21
+
22
+ def string_dtype_highest_priority(dtype1, dtype2):
23
+ if HAS_PYARROW:
24
+ DTYPE_HIERARCHY = [
25
+ StringDtype("python", na_value=np.nan),
26
+ StringDtype("pyarrow", na_value=np.nan),
27
+ StringDtype("python", na_value=NA),
28
+ StringDtype("pyarrow", na_value=NA),
29
+ ]
30
+ else:
31
+ DTYPE_HIERARCHY = [
32
+ StringDtype("python", na_value=np.nan),
33
+ StringDtype("python", na_value=NA),
34
+ ]
35
+
36
+ h1 = DTYPE_HIERARCHY.index(dtype1)
37
+ h2 = DTYPE_HIERARCHY.index(dtype2)
38
+ return DTYPE_HIERARCHY[max(h1, h2)]
39
+
40
+
41
+ def test_eq_all_na():
42
+ pytest.importorskip("pyarrow")
43
+ a = pd.array([NA, NA], dtype=StringDtype("pyarrow"))
44
+ result = a == a
45
+ expected = pd.array([NA, NA], dtype="boolean[pyarrow]")
46
+ tm.assert_extension_array_equal(result, expected)
47
+
48
+
49
+ def test_reversed_logical_ops(any_string_dtype):
50
+ # GH#60234
51
+ dtype = any_string_dtype
52
+ warn = None if dtype == object else Pandas4Warning
53
+ left = Series([True, False, False, True])
54
+ right = Series(["", "", "b", "c"], dtype=dtype)
55
+
56
+ msg = "operations between boolean dtype and"
57
+ with tm.assert_produces_warning(warn, match=msg):
58
+ result = left | right
59
+ expected = left | right.astype(bool)
60
+ tm.assert_series_equal(result, expected)
61
+
62
+ with tm.assert_produces_warning(warn, match=msg):
63
+ result = left & right
64
+ expected = left & right.astype(bool)
65
+ tm.assert_series_equal(result, expected)
66
+
67
+ with tm.assert_produces_warning(warn, match=msg):
68
+ result = left ^ right
69
+ expected = left ^ right.astype(bool)
70
+ tm.assert_series_equal(result, expected)
71
+
72
+
73
+ def test_pathlib_path_division(any_string_dtype, request):
74
+ # GH#61940
75
+ if any_string_dtype == object:
76
+ mark = pytest.mark.xfail(
77
+ reason="with NA present we go through _masked_arith_op which "
78
+ "raises TypeError bc Path is not recognized by lib.is_scalar."
79
+ )
80
+ request.applymarker(mark)
81
+
82
+ item = Path("/Users/Irv/")
83
+ ser = Series(["A", "B", NA], dtype=any_string_dtype)
84
+
85
+ result = item / ser
86
+ expected = Series([item / "A", item / "B", ser.dtype.na_value], dtype=object)
87
+ tm.assert_series_equal(result, expected)
88
+
89
+ result = ser / item
90
+ expected = Series(["A" / item, "B" / item, ser.dtype.na_value], dtype=object)
91
+ tm.assert_series_equal(result, expected)
92
+
93
+
94
+ def test_arithmetic_custom_object(any_string_dtype):
95
+ # GH#64107
96
+ class CustomObject:
97
+ def __init__(self, val):
98
+ self.val = val
99
+
100
+ def __add__(self, other):
101
+ if hasattr(other, "dtype"):
102
+ return NotImplemented
103
+ return CustomObject(self.val + other)
104
+
105
+ def __radd__(self, other):
106
+ if hasattr(other, "dtype"):
107
+ return NotImplemented
108
+ return CustomObject(other + self.val)
109
+
110
+ def __eq__(self, other):
111
+ return isinstance(other, CustomObject) and self.val == other.val
112
+
113
+ def __repr__(self):
114
+ return f"<CustomObject({self.val})>"
115
+
116
+ arr = np.array([CustomObject("a"), CustomObject("b")], dtype=object)
117
+ ser = Series(["1", "2"], dtype=any_string_dtype)
118
+ obj_ser = Series(arr)
119
+
120
+ cases = [
121
+ (ser + arr[0], [CustomObject("1a"), CustomObject("2a")]),
122
+ (ser + arr, [CustomObject("1a"), CustomObject("2b")]),
123
+ (ser + obj_ser, [CustomObject("1a"), CustomObject("2b")]),
124
+ (arr[0] + ser, [CustomObject("a1"), CustomObject("a2")]),
125
+ (arr + ser, [CustomObject("a1"), CustomObject("b2")]),
126
+ (obj_ser + ser, [CustomObject("a1"), CustomObject("b2")]),
127
+ ]
128
+ for result, expected_values in cases:
129
+ expected = Series(expected_values, dtype=object)
130
+ tm.assert_series_equal(result, expected)
131
+
132
+
133
+ def test_mixed_object_comparison(any_string_dtype):
134
+ # GH#60228
135
+ dtype = any_string_dtype
136
+ ser = Series(["a", "b"], dtype=dtype)
137
+
138
+ mixed = Series([1, "b"], dtype=object)
139
+
140
+ result = ser == mixed
141
+ expected = Series([False, True], dtype=bool)
142
+ if dtype == object:
143
+ pass
144
+ elif dtype.storage == "python" and dtype.na_value is NA:
145
+ expected = expected.astype("boolean")
146
+ elif dtype.storage == "pyarrow" and dtype.na_value is NA:
147
+ expected = expected.astype("bool[pyarrow]")
148
+
149
+ tm.assert_series_equal(result, expected)
150
+
151
+
152
+ def test_pyarrow_numpy_string_invalid():
153
+ # GH#56008
154
+ pa = pytest.importorskip("pyarrow")
155
+ ser = Series([False, True])
156
+ ser2 = Series(["a", "b"], dtype=StringDtype(na_value=np.nan))
157
+ result = ser == ser2
158
+ expected_eq = Series(False, index=ser.index)
159
+ tm.assert_series_equal(result, expected_eq)
160
+
161
+ result = ser != ser2
162
+ expected_ne = Series(True, index=ser.index)
163
+ tm.assert_series_equal(result, expected_ne)
164
+
165
+ with pytest.raises(TypeError, match="Invalid comparison"):
166
+ ser > ser2
167
+
168
+ # GH#59505
169
+ ser3 = ser2.astype("string[pyarrow]")
170
+ result3_eq = ser3 == ser
171
+ tm.assert_series_equal(result3_eq, expected_eq.astype("bool[pyarrow]"))
172
+ result3_ne = ser3 != ser
173
+ tm.assert_series_equal(result3_ne, expected_ne.astype("bool[pyarrow]"))
174
+
175
+ with pytest.raises(TypeError, match="Invalid comparison"):
176
+ ser > ser3
177
+
178
+ ser4 = ser2.astype(ArrowDtype(pa.string()))
179
+ result4_eq = ser4 == ser
180
+ tm.assert_series_equal(result4_eq, expected_eq.astype("bool[pyarrow]"))
181
+ result4_ne = ser4 != ser
182
+ tm.assert_series_equal(result4_ne, expected_ne.astype("bool[pyarrow]"))
183
+
184
+ with pytest.raises(TypeError, match="Invalid comparison"):
185
+ ser > ser4
186
+
187
+
188
+ def test_mul_bool_invalid(any_string_dtype):
189
+ # GH#62595
190
+ dtype = any_string_dtype
191
+ ser = Series(["a", "b", "c"], dtype=dtype)
192
+
193
+ if dtype == object:
194
+ pytest.skip("This is not expect to raise")
195
+ elif dtype.storage == "python":
196
+ msg = "Cannot multiply StringArray by bools. Explicitly cast to integers"
197
+ else:
198
+ msg = "Can only string multiply by an integer"
199
+
200
+ with pytest.raises(TypeError, match=msg):
201
+ False * ser
202
+ with pytest.raises(TypeError, match=msg):
203
+ ser * True
204
+ with pytest.raises(TypeError, match=msg):
205
+ ser * np.array([True, False, True], dtype=bool)
206
+ with pytest.raises(TypeError, match=msg):
207
+ np.array([True, False, True], dtype=bool) * ser
208
+
209
+
210
+ def test_add(any_string_dtype, request):
211
+ dtype = any_string_dtype
212
+ if dtype == object:
213
+ mark = pytest.mark.xfail(
214
+ reason="Need to update expected for numpy object dtype"
215
+ )
216
+ request.applymarker(mark)
217
+
218
+ a = Series(["a", "b", "c", None, None], dtype=dtype)
219
+ b = Series(["x", "y", None, "z", None], dtype=dtype)
220
+
221
+ result = a + b
222
+ expected = Series(["ax", "by", None, None, None], dtype=dtype)
223
+ tm.assert_series_equal(result, expected)
224
+
225
+ result = a.add(b)
226
+ tm.assert_series_equal(result, expected)
227
+
228
+ result = a.radd(b)
229
+ expected = Series(["xa", "yb", None, None, None], dtype=dtype)
230
+ tm.assert_series_equal(result, expected)
231
+
232
+ result = a.add(b, fill_value="-")
233
+ expected = Series(["ax", "by", "c-", "-z", None], dtype=dtype)
234
+ tm.assert_series_equal(result, expected)
235
+
236
+
237
+ def test_add_2d(any_string_dtype, request):
238
+ dtype = any_string_dtype
239
+
240
+ if dtype == object:
241
+ reason = "Failed: DID NOT RAISE <class 'ValueError'>"
242
+ mark = pytest.mark.xfail(raises=None, reason=reason)
243
+ request.applymarker(mark)
244
+
245
+ a = pd.array(["a", "b", "c"], dtype=dtype)
246
+ b = np.array([["a", "b", "c"]], dtype=object)
247
+ with pytest.raises(ValueError, match="3 != 1"):
248
+ a + b
249
+
250
+ s = Series(a)
251
+ with pytest.raises(ValueError, match="3 != 1"):
252
+ s + b
253
+
254
+
255
+ def test_add_sequence(any_string_dtype, request, using_infer_string):
256
+ dtype = any_string_dtype
257
+ if (
258
+ dtype != object
259
+ and dtype.storage == "python"
260
+ and dtype.na_value is np.nan
261
+ and HAS_PYARROW
262
+ and using_infer_string
263
+ ):
264
+ mark = pytest.mark.xfail(
265
+ reason="As of GH#62522, the list gets wrapped with sanitize_array, "
266
+ "which casts to a higher-priority StringArray, so we get "
267
+ "NotImplemented."
268
+ )
269
+ request.applymarker(mark)
270
+ if dtype == np.dtype(object) and using_infer_string:
271
+ mark = pytest.mark.xfail(reason="Cannot broadcast list")
272
+ request.applymarker(mark)
273
+
274
+ a = pd.array(["a", "b", None, None], dtype=dtype)
275
+ other = ["x", None, "y", None]
276
+
277
+ result = a + other
278
+ expected = pd.array(["ax", None, None, None], dtype=dtype)
279
+ tm.assert_extension_array_equal(result, expected)
280
+
281
+ result = other + a
282
+ expected = pd.array(["xa", None, None, None], dtype=dtype)
283
+ tm.assert_extension_array_equal(result, expected)
284
+
285
+
286
+ def test_string_add_missing_values(string_dtype_no_object):
287
+ # GH#64968 Arrow-backed str arrays should return NA when added to missing
288
+ arr = pd.array(["y"], dtype=string_dtype_no_object)
289
+ expected = pd.array([NA], dtype=string_dtype_no_object)
290
+
291
+ for na_val in [None, np.nan, NA]:
292
+ # left side
293
+ result = arr + na_val
294
+ tm.assert_extension_array_equal(result, expected)
295
+
296
+ # right side
297
+ result = na_val + arr
298
+ tm.assert_extension_array_equal(result, expected)
299
+
300
+
301
+ def test_mul(any_string_dtype):
302
+ dtype = any_string_dtype
303
+ a = pd.array(["a", "b", None], dtype=dtype)
304
+ result = a * 2
305
+ expected = pd.array(["aa", "bb", None], dtype=dtype)
306
+ tm.assert_extension_array_equal(result, expected)
307
+
308
+ result = 2 * a
309
+ tm.assert_extension_array_equal(result, expected)
310
+
311
+
312
+ def test_add_strings(any_string_dtype, request):
313
+ dtype = any_string_dtype
314
+ if dtype != np.dtype(object):
315
+ mark = pytest.mark.xfail(reason="GH-28527")
316
+ request.applymarker(mark)
317
+ arr = pd.array(["a", "b", "c", "d"], dtype=dtype)
318
+ df = pd.DataFrame([["t", "y", "v", "w"]], dtype=object)
319
+ assert arr.__add__(df) is NotImplemented
320
+
321
+ result = arr + df
322
+ expected = pd.DataFrame([["at", "by", "cv", "dw"]]).astype(dtype)
323
+ tm.assert_frame_equal(result, expected)
324
+
325
+ result = df + arr
326
+ expected = pd.DataFrame([["ta", "yb", "vc", "wd"]]).astype(dtype)
327
+ tm.assert_frame_equal(result, expected)
328
+
329
+
330
+ @pytest.mark.xfail(reason="GH-28527")
331
+ def test_add_frame(dtype):
332
+ arr = pd.array(["a", "b", np.nan, np.nan], dtype=dtype)
333
+ df = pd.DataFrame([["x", np.nan, "y", np.nan]])
334
+
335
+ assert arr.__add__(df) is NotImplemented
336
+
337
+ result = arr + df
338
+ expected = pd.DataFrame([["ax", np.nan, np.nan, np.nan]]).astype(dtype)
339
+ tm.assert_frame_equal(result, expected)
340
+
341
+ result = df + arr
342
+ expected = pd.DataFrame([["xa", np.nan, np.nan, np.nan]]).astype(dtype)
343
+ tm.assert_frame_equal(result, expected)
344
+
345
+
346
+ def test_comparison_methods_scalar(comparison_op, any_string_dtype):
347
+ dtype = any_string_dtype
348
+ op_name = f"__{comparison_op.__name__}__"
349
+ a = pd.array(["a", None, "c"], dtype=dtype)
350
+ other = "a"
351
+ result = getattr(a, op_name)(other)
352
+ if dtype == object or dtype.na_value is np.nan:
353
+ expected = np.array([getattr(item, op_name)(other) for item in a])
354
+ if comparison_op == operator.ne:
355
+ expected[1] = True
356
+ else:
357
+ expected[1] = False
358
+ result = extract_array(result, extract_numpy=True)
359
+ tm.assert_numpy_array_equal(result, expected.astype(np.bool_))
360
+ else:
361
+ expected_dtype = "boolean[pyarrow]" if dtype.storage == "pyarrow" else "boolean"
362
+ expected = np.array([getattr(item, op_name)(other) for item in a], dtype=object)
363
+ expected = pd.array(expected, dtype=expected_dtype)
364
+ tm.assert_extension_array_equal(result, expected)
365
+
366
+
367
+ def test_comparison_methods_scalar_pd_na(comparison_op, any_string_dtype):
368
+ dtype = any_string_dtype
369
+ op_name = f"__{comparison_op.__name__}__"
370
+ a = pd.array(["a", None, "c"], dtype=dtype)
371
+ result = getattr(a, op_name)(NA)
372
+
373
+ if dtype == np.dtype(object) or dtype.na_value is np.nan:
374
+ if operator.ne == comparison_op:
375
+ expected = np.array([True, True, True])
376
+ else:
377
+ expected = np.array([False, False, False])
378
+ result = extract_array(result, extract_numpy=True)
379
+ tm.assert_numpy_array_equal(result, expected)
380
+ else:
381
+ expected_dtype = "boolean[pyarrow]" if dtype.storage == "pyarrow" else "boolean"
382
+ expected = pd.array([None, None, None], dtype=expected_dtype)
383
+ tm.assert_extension_array_equal(result, expected)
384
+ tm.assert_extension_array_equal(result, expected)
385
+
386
+
387
+ def test_comparison_methods_scalar_not_string(comparison_op, any_string_dtype):
388
+ op_name = f"__{comparison_op.__name__}__"
389
+ dtype = any_string_dtype
390
+
391
+ a = pd.array(["a", None, "c"], dtype=dtype)
392
+ other = 42
393
+
394
+ if op_name not in ["__eq__", "__ne__"]:
395
+ with pytest.raises(TypeError, match="Invalid comparison|not supported between"):
396
+ getattr(a, op_name)(other)
397
+
398
+ return
399
+
400
+ result = getattr(a, op_name)(other)
401
+ result = extract_array(result, extract_numpy=True)
402
+
403
+ if dtype == np.dtype(object) or dtype.na_value is np.nan:
404
+ expected_data = {
405
+ "__eq__": [False, False, False],
406
+ "__ne__": [True, True, True],
407
+ }[op_name]
408
+ expected = np.array(expected_data)
409
+ tm.assert_numpy_array_equal(result, expected)
410
+ else:
411
+ expected_data = {"__eq__": [False, None, False], "__ne__": [True, None, True]}[
412
+ op_name
413
+ ]
414
+ expected_dtype = "boolean[pyarrow]" if dtype.storage == "pyarrow" else "boolean"
415
+ expected = pd.array(expected_data, dtype=expected_dtype)
416
+ tm.assert_extension_array_equal(result, expected)
417
+
418
+
419
+ def test_comparison_methods_array(comparison_op, any_string_dtype, any_string_dtype2):
420
+ op_name = f"__{comparison_op.__name__}__"
421
+ dtype = any_string_dtype
422
+ dtype2 = any_string_dtype2
423
+
424
+ a = pd.array(["a", None, "c"], dtype=dtype)
425
+ other = pd.array([None, None, "c"], dtype=dtype2)
426
+ result = comparison_op(a, other)
427
+ result = extract_array(result, extract_numpy=True)
428
+
429
+ # ensure operation is commutative
430
+ result2 = comparison_op(other, a)
431
+ result2 = extract_array(result2, extract_numpy=True)
432
+ tm.assert_equal(result, result2)
433
+
434
+ if (dtype == object or dtype.na_value is np.nan) and (
435
+ dtype2 == object or dtype2.na_value is np.nan
436
+ ):
437
+ if operator.ne == comparison_op:
438
+ expected = np.array([True, True, False])
439
+ else:
440
+ expected = np.array([False, False, False])
441
+ expected[-1] = getattr(other[-1], op_name)(a[-1])
442
+ result = extract_array(result, extract_numpy=True)
443
+ tm.assert_numpy_array_equal(result, expected)
444
+
445
+ else:
446
+ if dtype == object:
447
+ max_dtype = dtype2
448
+ elif dtype2 == object:
449
+ max_dtype = dtype
450
+ else:
451
+ max_dtype = string_dtype_highest_priority(dtype, dtype2)
452
+ if max_dtype.storage == "python":
453
+ expected_dtype = "boolean"
454
+ else:
455
+ expected_dtype = "bool[pyarrow]"
456
+
457
+ expected = np.full(len(a), fill_value=None, dtype="object")
458
+ expected[-1] = getattr(other[-1], op_name)(a[-1])
459
+ expected = pd.array(expected, dtype=expected_dtype)
460
+ tm.assert_equal(result, expected)
461
+
462
+
463
+ @td.skip_if_no("pyarrow")
464
+ def test_comparison_methods_array_arrow_extension(comparison_op, any_string_dtype):
465
+ # Test pd.ArrowDtype(pa.string()) against other string arrays
466
+ import pyarrow as pa
467
+
468
+ dtype2 = any_string_dtype
469
+
470
+ op_name = f"__{comparison_op.__name__}__"
471
+ dtype = ArrowDtype(pa.string())
472
+ a = pd.array(["a", None, "c"], dtype=dtype)
473
+ other = pd.array([None, None, "c"], dtype=dtype2)
474
+ result = comparison_op(a, other)
475
+
476
+ # ensure operation is commutative
477
+ result2 = comparison_op(other, a)
478
+ tm.assert_equal(result, result2)
479
+
480
+ expected = pd.array([None, None, True], dtype="bool[pyarrow]")
481
+ expected[-1] = getattr(other[-1], op_name)(a[-1])
482
+ tm.assert_extension_array_equal(result, expected)
483
+
484
+
485
+ @pytest.mark.parametrize("box", [pd.array, pd.Index, Series])
486
+ def test_comparison_methods_list(comparison_op, any_string_dtype, box, request):
487
+ dtype = any_string_dtype
488
+
489
+ if box is pd.array and dtype != object and dtype.na_value is np.nan:
490
+ mark = pytest.mark.xfail(
491
+ reason="After wrapping list, op returns NotImplemented, see GH#62522"
492
+ )
493
+ request.applymarker(mark)
494
+
495
+ op_name = f"__{comparison_op.__name__}__"
496
+
497
+ a = box(pd.array(["a", None, "c"], dtype=dtype))
498
+ item = "c"
499
+ other = [None, None, "c"]
500
+ result = comparison_op(a, other)
501
+
502
+ # ensure operation is commutative
503
+ result2 = comparison_op(other, a)
504
+ tm.assert_equal(result, result2)
505
+
506
+ if dtype == np.dtype(object) or dtype.na_value is np.nan:
507
+ if operator.ne == comparison_op:
508
+ expected = np.array([True, True, False])
509
+ else:
510
+ expected = np.array([False, False, False])
511
+ expected[-1] = getattr(item, op_name)(item)
512
+ if box is not pd.Index:
513
+ # if GH#62766 is addressed this check can be removed
514
+ expected = box(expected, dtype=expected.dtype)
515
+ tm.assert_equal(result, expected)
516
+
517
+ else:
518
+ expected_dtype = "boolean[pyarrow]" if dtype.storage == "pyarrow" else "boolean"
519
+ expected = np.full(len(a), fill_value=None, dtype="object")
520
+ expected[-1] = getattr(item, op_name)(item)
521
+ expected = pd.array(expected, dtype=expected_dtype)
522
+ expected = extract_array(expected, extract_numpy=True)
523
+ if box is not pd.Index:
524
+ # if GH#62766 is addressed this check can be removed
525
+ expected = tm.box_expected(expected, box)
526
+ tm.assert_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arithmetic/test_timedelta64.py ADDED
@@ -0,0 +1,2336 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arithmetic tests for DataFrame/Series/Index/Array classes that should
2
+ # behave identically.
3
+ from datetime import (
4
+ datetime,
5
+ timedelta,
6
+ )
7
+
8
+ import numpy as np
9
+ import pytest
10
+
11
+ from pandas._libs.tslibs import timezones
12
+ from pandas.compat import WASM
13
+ from pandas.errors import OutOfBoundsDatetime
14
+ import pandas.util._test_decorators as td
15
+
16
+ import pandas as pd
17
+ from pandas import (
18
+ DataFrame,
19
+ DatetimeIndex,
20
+ Index,
21
+ NaT,
22
+ Series,
23
+ Timedelta,
24
+ TimedeltaIndex,
25
+ Timestamp,
26
+ offsets,
27
+ timedelta_range,
28
+ )
29
+ import pandas._testing as tm
30
+ from pandas.core.arrays import NumpyExtensionArray
31
+ from pandas.tests.arithmetic.common import (
32
+ assert_invalid_addsub_type,
33
+ assert_invalid_comparison,
34
+ get_upcast_box,
35
+ )
36
+
37
+
38
+ def assert_dtype(obj, expected_dtype):
39
+ """
40
+ Helper to check the dtype for a Series, Index, or single-column DataFrame.
41
+ """
42
+ dtype = tm.get_dtype(obj)
43
+
44
+ assert dtype == expected_dtype
45
+
46
+
47
+ def get_expected_name(box, names):
48
+ if box is DataFrame:
49
+ # Since we are operating with a DataFrame and a non-DataFrame,
50
+ # the non-DataFrame is cast to Series and its name ignored.
51
+ exname = names[0]
52
+ elif box in [tm.to_array, pd.array]:
53
+ exname = names[1]
54
+ else:
55
+ exname = names[2]
56
+ return exname
57
+
58
+
59
+ # ------------------------------------------------------------------
60
+ # Timedelta64[ns] dtype Comparisons
61
+
62
+
63
+ class TestTimedelta64ArrayLikeComparisons:
64
+ # Comparison tests for timedelta64[ns] vectors fully parametrized over
65
+ # DataFrame/Series/TimedeltaIndex/TimedeltaArray. Ideally all comparison
66
+ # tests will eventually end up here.
67
+
68
+ def test_compare_timedelta64_zerodim(self, box_with_array):
69
+ # GH#26689 should unbox when comparing with zerodim array
70
+ box = box_with_array
71
+ xbox = box_with_array if box_with_array not in [Index, pd.array] else np.ndarray
72
+
73
+ tdi = timedelta_range("2h", periods=4)
74
+ other = np.array(tdi.to_numpy()[0])
75
+
76
+ tdi = tm.box_expected(tdi, box)
77
+ res = tdi <= other
78
+ expected = np.array([True, False, False, False])
79
+ expected = tm.box_expected(expected, xbox)
80
+ tm.assert_equal(res, expected)
81
+
82
+ @pytest.mark.parametrize(
83
+ "td_scalar",
84
+ [
85
+ timedelta(days=1),
86
+ Timedelta(days=1),
87
+ Timedelta(days=1).to_timedelta64(),
88
+ offsets.Hour(24),
89
+ ],
90
+ )
91
+ def test_compare_timedeltalike_scalar(self, box_with_array, td_scalar):
92
+ # regression test for GH#5963
93
+ box = box_with_array
94
+ xbox = box if box not in [Index, pd.array] else np.ndarray
95
+
96
+ ser = Series([timedelta(days=1), timedelta(days=2)])
97
+ ser = tm.box_expected(ser, box)
98
+ actual = ser > td_scalar
99
+ expected = Series([False, True])
100
+ expected = tm.box_expected(expected, xbox)
101
+ tm.assert_equal(actual, expected)
102
+
103
+ @pytest.mark.parametrize(
104
+ "invalid",
105
+ [
106
+ 345600000000000,
107
+ "a",
108
+ Timestamp("2021-01-01"),
109
+ Timestamp("2021-01-01").now("UTC"),
110
+ Timestamp("2021-01-01").now().to_datetime64(),
111
+ Timestamp("2021-01-01").now().to_pydatetime(),
112
+ Timestamp("2021-01-01").date(),
113
+ np.array(4), # zero-dim mismatched dtype
114
+ ],
115
+ )
116
+ def test_td64_comparisons_invalid(self, box_with_array, invalid):
117
+ # GH#13624 for str
118
+ box = box_with_array
119
+
120
+ rng = timedelta_range("1 days", periods=10)
121
+ obj = tm.box_expected(rng, box)
122
+
123
+ assert_invalid_comparison(obj, invalid, box)
124
+
125
+ @pytest.mark.parametrize(
126
+ "other",
127
+ [
128
+ list(range(10)),
129
+ np.arange(10),
130
+ np.arange(10).astype(np.float32),
131
+ np.arange(10).astype(object),
132
+ pd.date_range("1970-01-01", periods=10, tz="UTC").array,
133
+ np.array(pd.date_range("1970-01-01", periods=10)),
134
+ list(pd.date_range("1970-01-01", periods=10)),
135
+ pd.date_range("1970-01-01", periods=10).astype(object),
136
+ pd.period_range("1971-01-01", freq="D", periods=10).array,
137
+ pd.period_range("1971-01-01", freq="D", periods=10).astype(object),
138
+ ],
139
+ )
140
+ def test_td64arr_cmp_arraylike_invalid(self, other, box_with_array):
141
+ # We don't parametrize this over box_with_array because listlike
142
+ # other plays poorly with assert_invalid_comparison reversed checks
143
+
144
+ rng = timedelta_range("1 days", periods=10)._data
145
+ rng = tm.box_expected(rng, box_with_array)
146
+ assert_invalid_comparison(rng, other, box_with_array)
147
+
148
+ def test_td64arr_cmp_mixed_invalid(self):
149
+ rng = timedelta_range("1 days", periods=5)._data
150
+ other = np.array([0, 1, 2, rng[3], Timestamp("2021-01-01")])
151
+
152
+ result = rng == other
153
+ expected = np.array([False, False, False, True, False])
154
+ tm.assert_numpy_array_equal(result, expected)
155
+
156
+ result = rng != other
157
+ tm.assert_numpy_array_equal(result, ~expected)
158
+
159
+ msg = "Invalid comparison between|Cannot compare type|not supported between"
160
+ with pytest.raises(TypeError, match=msg):
161
+ rng < other
162
+ with pytest.raises(TypeError, match=msg):
163
+ rng > other
164
+ with pytest.raises(TypeError, match=msg):
165
+ rng <= other
166
+ with pytest.raises(TypeError, match=msg):
167
+ rng >= other
168
+
169
+
170
+ class TestTimedelta64ArrayComparisons:
171
+ # TODO: All of these need to be parametrized over box
172
+
173
+ @pytest.mark.parametrize("dtype", [None, object])
174
+ def test_comp_nat(self, dtype):
175
+ left = TimedeltaIndex([Timedelta("1 days"), NaT, Timedelta("3 days")])
176
+ right = TimedeltaIndex([NaT, NaT, Timedelta("3 days")])
177
+
178
+ lhs, rhs = left, right
179
+ if dtype is object:
180
+ lhs, rhs = left.astype(object), right.astype(object)
181
+
182
+ result = rhs == lhs
183
+ expected = np.array([False, False, True])
184
+ tm.assert_numpy_array_equal(result, expected)
185
+
186
+ result = rhs != lhs
187
+ expected = np.array([True, True, False])
188
+ tm.assert_numpy_array_equal(result, expected)
189
+
190
+ expected = np.array([False, False, False])
191
+ tm.assert_numpy_array_equal(lhs == NaT, expected)
192
+ tm.assert_numpy_array_equal(NaT == rhs, expected)
193
+
194
+ expected = np.array([True, True, True])
195
+ tm.assert_numpy_array_equal(lhs != NaT, expected)
196
+ tm.assert_numpy_array_equal(NaT != lhs, expected)
197
+
198
+ expected = np.array([False, False, False])
199
+ tm.assert_numpy_array_equal(lhs < NaT, expected)
200
+ tm.assert_numpy_array_equal(NaT > lhs, expected)
201
+
202
+ @pytest.mark.parametrize(
203
+ "idx2",
204
+ [
205
+ TimedeltaIndex(
206
+ ["2 day", "2 day", NaT, NaT, "1 day 00:00:02", "5 days 00:00:03"]
207
+ ),
208
+ np.array(
209
+ [
210
+ np.timedelta64(2, "D"),
211
+ np.timedelta64(2, "D"),
212
+ np.timedelta64("NaT", "ns"),
213
+ np.timedelta64("NaT", "ns"),
214
+ np.timedelta64(1, "D") + np.timedelta64(2, "s"),
215
+ np.timedelta64(5, "D") + np.timedelta64(3, "s"),
216
+ ]
217
+ ),
218
+ ],
219
+ )
220
+ def test_comparisons_nat(self, idx2):
221
+ idx1 = TimedeltaIndex(
222
+ [
223
+ "1 day",
224
+ NaT,
225
+ "1 day 00:00:01",
226
+ NaT,
227
+ "1 day 00:00:01",
228
+ "5 day 00:00:03",
229
+ ]
230
+ )
231
+ # Check pd.NaT is handles as the same as np.nan
232
+ result = idx1 < idx2
233
+ expected = np.array([True, False, False, False, True, False])
234
+ tm.assert_numpy_array_equal(result, expected)
235
+
236
+ result = idx2 > idx1
237
+ expected = np.array([True, False, False, False, True, False])
238
+ tm.assert_numpy_array_equal(result, expected)
239
+
240
+ result = idx1 <= idx2
241
+ expected = np.array([True, False, False, False, True, True])
242
+ tm.assert_numpy_array_equal(result, expected)
243
+
244
+ result = idx2 >= idx1
245
+ expected = np.array([True, False, False, False, True, True])
246
+ tm.assert_numpy_array_equal(result, expected)
247
+
248
+ result = idx1 == idx2
249
+ expected = np.array([False, False, False, False, False, True])
250
+ tm.assert_numpy_array_equal(result, expected)
251
+
252
+ result = idx1 != idx2
253
+ expected = np.array([True, True, True, True, True, False])
254
+ tm.assert_numpy_array_equal(result, expected)
255
+
256
+ # TODO: better name
257
+ def test_comparisons_coverage(self):
258
+ rng = timedelta_range("1 days", periods=10)
259
+
260
+ result = rng < rng[3]
261
+ expected = np.array([True, True, True] + [False] * 7)
262
+ tm.assert_numpy_array_equal(result, expected)
263
+
264
+ result = rng == list(rng)
265
+ exp = rng == rng
266
+ tm.assert_numpy_array_equal(result, exp)
267
+
268
+
269
+ # ------------------------------------------------------------------
270
+ # Timedelta64[ns] dtype Arithmetic Operations
271
+
272
+
273
+ class TestTimedelta64ArithmeticUnsorted:
274
+ # Tests moved from type-specific test files but not
275
+ # yet sorted/parametrized/de-duplicated
276
+
277
+ def test_td64_op_with_list(self, box_with_array):
278
+ # GH#62353
279
+ box = box_with_array
280
+
281
+ left = TimedeltaIndex(["2D", "4D"])
282
+ left = tm.box_expected(left, box)
283
+
284
+ right = [Timestamp("2016-01-01"), Timestamp("2016-02-01")]
285
+
286
+ result = left + right
287
+ expected = DatetimeIndex(["2016-01-03", "2016-02-05"], dtype="M8[us]")
288
+ expected = tm.box_expected(expected, box)
289
+ tm.assert_equal(result, expected)
290
+
291
+ result2 = right + left
292
+ tm.assert_equal(result2, expected)
293
+
294
+ def test_ufunc_coercions(self):
295
+ # normal ops are also tested in tseries/test_timedeltas.py
296
+ idx = TimedeltaIndex(["2h", "4h", "6h", "8h", "10h"], freq="2h", name="x")
297
+
298
+ for result in [idx * 2, np.multiply(idx, 2)]:
299
+ assert isinstance(result, TimedeltaIndex)
300
+ exp = TimedeltaIndex(["4h", "8h", "12h", "16h", "20h"], freq="4h", name="x")
301
+ tm.assert_index_equal(result, exp)
302
+ assert result.freq == "4h"
303
+
304
+ for result in [idx / 2, np.divide(idx, 2)]:
305
+ assert isinstance(result, TimedeltaIndex)
306
+ exp = TimedeltaIndex(["1h", "2h", "3h", "4h", "5h"], freq="h", name="x")
307
+ tm.assert_index_equal(result, exp)
308
+ assert result.freq == "h"
309
+
310
+ for result in [-idx, np.negative(idx)]:
311
+ assert isinstance(result, TimedeltaIndex)
312
+ exp = TimedeltaIndex(
313
+ ["-2h", "-4h", "-6h", "-8h", "-10h"], freq="-2h", name="x"
314
+ )
315
+ tm.assert_index_equal(result, exp)
316
+ assert result.freq == "-2h"
317
+
318
+ idx = TimedeltaIndex(["-2h", "-1h", "0h", "1h", "2h"], freq="h", name="x")
319
+ for result in [abs(idx), np.absolute(idx)]:
320
+ assert isinstance(result, TimedeltaIndex)
321
+ exp = TimedeltaIndex(["2h", "1h", "0h", "1h", "2h"], freq=None, name="x")
322
+ tm.assert_index_equal(result, exp)
323
+ assert result.freq is None
324
+
325
+ def test_subtraction_ops(self):
326
+ # with datetimes/timedelta and tdi/dti
327
+ tdi = TimedeltaIndex(["1 days", NaT, "2 days"], name="foo")
328
+ dti = pd.date_range("20130101", periods=3, name="bar")
329
+ td = Timedelta("1 days")
330
+ dt = Timestamp("20130101")
331
+
332
+ msg = "cannot subtract a datelike from a TimedeltaArray"
333
+ with pytest.raises(TypeError, match=msg):
334
+ tdi - dt
335
+ with pytest.raises(TypeError, match=msg):
336
+ tdi - dti
337
+
338
+ msg = r"unsupported operand type\(s\) for -"
339
+ with pytest.raises(TypeError, match=msg):
340
+ td - dt
341
+
342
+ msg = "cannot subtract DatetimeArray from Timedelta"
343
+ with pytest.raises(TypeError, match=msg):
344
+ td - dti
345
+
346
+ result = dt - dti
347
+ expected = TimedeltaIndex(["0 days", "-1 days", "-2 days"], name="bar")
348
+ tm.assert_index_equal(result, expected)
349
+
350
+ result = dti - dt
351
+ expected = TimedeltaIndex(["0 days", "1 days", "2 days"], name="bar")
352
+ tm.assert_index_equal(result, expected)
353
+
354
+ result = tdi - td
355
+ expected = TimedeltaIndex(["0 days", NaT, "1 days"], name="foo")
356
+ tm.assert_index_equal(result, expected)
357
+
358
+ result = td - tdi
359
+ expected = TimedeltaIndex(["0 days", NaT, "-1 days"], name="foo")
360
+ tm.assert_index_equal(result, expected)
361
+
362
+ result = dti - td
363
+ expected = DatetimeIndex(
364
+ ["20121231", "20130101", "20130102"], dtype="M8[us]", freq="D", name="bar"
365
+ )
366
+ tm.assert_index_equal(result, expected)
367
+
368
+ result = dt - tdi
369
+ expected = DatetimeIndex(
370
+ ["20121231", NaT, "20121230"], dtype="M8[us]", name="foo"
371
+ )
372
+ tm.assert_index_equal(result, expected)
373
+
374
+ def test_subtraction_ops_with_tz(self, box_with_array):
375
+ # check that dt/dti subtraction ops with tz are validated
376
+ dti = pd.date_range("20130101", periods=3)
377
+ dti = tm.box_expected(dti, box_with_array)
378
+ ts = Timestamp("20130101")
379
+ dt = ts.to_pydatetime()
380
+ dti_tz = pd.date_range("20130101", periods=3).tz_localize("US/Eastern")
381
+ dti_tz = tm.box_expected(dti_tz, box_with_array)
382
+ ts_tz = Timestamp("20130101").tz_localize("US/Eastern")
383
+ ts_tz2 = Timestamp("20130101").tz_localize("CET")
384
+ dt_tz = ts_tz.to_pydatetime()
385
+ td = Timedelta("1 days")
386
+
387
+ def _check(result, expected):
388
+ assert result == expected
389
+ assert isinstance(result, Timedelta)
390
+
391
+ # scalars
392
+ result = ts - ts
393
+ expected = Timedelta("0 days")
394
+ _check(result, expected)
395
+
396
+ result = dt_tz - ts_tz
397
+ expected = Timedelta("0 days")
398
+ _check(result, expected)
399
+
400
+ result = ts_tz - dt_tz
401
+ expected = Timedelta("0 days")
402
+ _check(result, expected)
403
+
404
+ # tz mismatches
405
+ msg = "Cannot subtract tz-naive and tz-aware datetime-like objects."
406
+ with pytest.raises(TypeError, match=msg):
407
+ dt_tz - ts
408
+ msg = "can't subtract offset-naive and offset-aware datetimes"
409
+ with pytest.raises(TypeError, match=msg):
410
+ dt_tz - dt
411
+ msg = "can't subtract offset-naive and offset-aware datetimes"
412
+ with pytest.raises(TypeError, match=msg):
413
+ dt - dt_tz
414
+ msg = "Cannot subtract tz-naive and tz-aware datetime-like objects."
415
+ with pytest.raises(TypeError, match=msg):
416
+ ts - dt_tz
417
+ with pytest.raises(TypeError, match=msg):
418
+ ts_tz2 - ts
419
+ with pytest.raises(TypeError, match=msg):
420
+ ts_tz2 - dt
421
+
422
+ msg = "Cannot subtract tz-naive and tz-aware"
423
+ # with dti
424
+ with pytest.raises(TypeError, match=msg):
425
+ dti - ts_tz
426
+ with pytest.raises(TypeError, match=msg):
427
+ dti_tz - ts
428
+
429
+ result = dti_tz - dt_tz
430
+ expected = TimedeltaIndex(["0 days", "1 days", "2 days"])
431
+ expected = tm.box_expected(expected, box_with_array)
432
+ tm.assert_equal(result, expected)
433
+
434
+ result = dt_tz - dti_tz
435
+ expected = TimedeltaIndex(["0 days", "-1 days", "-2 days"])
436
+ expected = tm.box_expected(expected, box_with_array)
437
+ tm.assert_equal(result, expected)
438
+
439
+ result = dti_tz - ts_tz
440
+ expected = TimedeltaIndex(["0 days", "1 days", "2 days"])
441
+ expected = tm.box_expected(expected, box_with_array)
442
+ tm.assert_equal(result, expected)
443
+
444
+ result = ts_tz - dti_tz
445
+ expected = TimedeltaIndex(["0 days", "-1 days", "-2 days"])
446
+ expected = tm.box_expected(expected, box_with_array)
447
+ tm.assert_equal(result, expected)
448
+
449
+ result = td - td
450
+ expected = Timedelta("0 days")
451
+ _check(result, expected)
452
+
453
+ result = dti_tz - td
454
+ expected = DatetimeIndex(["20121231", "20130101", "20130102"], tz="US/Eastern")
455
+ expected = tm.box_expected(expected, box_with_array)
456
+ tm.assert_equal(result, expected)
457
+
458
+ def test_dti_tdi_numeric_ops(self):
459
+ # These are normally union/diff set-like ops
460
+ tdi = TimedeltaIndex(["1 days", NaT, "2 days"], name="foo")
461
+ dti = pd.date_range("20130101", periods=3, name="bar")
462
+
463
+ result = tdi - tdi
464
+ expected = TimedeltaIndex(["0 days", NaT, "0 days"], name="foo")
465
+ tm.assert_index_equal(result, expected)
466
+
467
+ result = tdi + tdi
468
+ expected = TimedeltaIndex(["2 days", NaT, "4 days"], name="foo")
469
+ tm.assert_index_equal(result, expected)
470
+
471
+ result = dti - tdi # name will be reset
472
+ expected = DatetimeIndex(["20121231", NaT, "20130101"], dtype="M8[us]")
473
+ tm.assert_index_equal(result, expected)
474
+
475
+ def test_addition_ops(self):
476
+ # with datetimes/timedelta and tdi/dti
477
+ tdi = TimedeltaIndex(["1 days", NaT, "2 days"], name="foo")
478
+ dti = pd.date_range("20130101", periods=3, name="bar")
479
+ td = Timedelta("1 days")
480
+ dt = Timestamp("20130101")
481
+
482
+ result = tdi + dt
483
+ expected = DatetimeIndex(
484
+ ["20130102", NaT, "20130103"], dtype="M8[us]", name="foo"
485
+ )
486
+ tm.assert_index_equal(result, expected)
487
+
488
+ result = dt + tdi
489
+ expected = DatetimeIndex(
490
+ ["20130102", NaT, "20130103"], dtype="M8[us]", name="foo"
491
+ )
492
+ tm.assert_index_equal(result, expected)
493
+
494
+ result = td + tdi
495
+ expected = TimedeltaIndex(["2 days", NaT, "3 days"], name="foo")
496
+ tm.assert_index_equal(result, expected)
497
+
498
+ result = tdi + td
499
+ expected = TimedeltaIndex(["2 days", NaT, "3 days"], name="foo")
500
+ tm.assert_index_equal(result, expected)
501
+
502
+ # unequal length
503
+ msg = "cannot add indices of unequal length"
504
+ with pytest.raises(ValueError, match=msg):
505
+ tdi + dti[0:1]
506
+ with pytest.raises(ValueError, match=msg):
507
+ tdi[0:1] + dti
508
+
509
+ # random indexes
510
+ msg = "Addition/subtraction of integers and integer-arrays"
511
+ with pytest.raises(TypeError, match=msg):
512
+ tdi + Index([1, 2, 3], dtype=np.int64)
513
+
514
+ # this is a union!
515
+ # FIXME: don't leave commented-out
516
+ # pytest.raises(TypeError, lambda : Index([1,2,3]) + tdi)
517
+
518
+ result = tdi + dti # name will be reset
519
+ expected = DatetimeIndex(["20130102", NaT, "20130105"], dtype="M8[us]")
520
+ tm.assert_index_equal(result, expected)
521
+
522
+ result = dti + tdi # name will be reset
523
+ expected = DatetimeIndex(["20130102", NaT, "20130105"], dtype="M8[us]")
524
+ tm.assert_index_equal(result, expected)
525
+
526
+ result = dt + td
527
+ expected = Timestamp("20130102")
528
+ assert result == expected
529
+
530
+ result = td + dt
531
+ expected = Timestamp("20130102")
532
+ assert result == expected
533
+
534
+ # TODO: Needs more informative name, probably split up into
535
+ # more targeted tests
536
+ @pytest.mark.parametrize("freq", ["D", "B"])
537
+ def test_timedelta(self, freq):
538
+ index = pd.date_range("1/1/2000", periods=50, freq=freq, unit="ns")
539
+
540
+ shifted = index + timedelta(1)
541
+ back = shifted + timedelta(-1)
542
+ back = back._with_freq("infer")
543
+ tm.assert_index_equal(index, back)
544
+
545
+ if freq == "D":
546
+ expected = pd.tseries.offsets.Day(1)
547
+ assert index.freq == expected
548
+ assert shifted.freq == expected
549
+ assert back.freq == expected
550
+ else: # freq == 'B'
551
+ assert index.freq == pd.tseries.offsets.BusinessDay(1)
552
+ assert shifted.freq is None
553
+ assert back.freq == pd.tseries.offsets.BusinessDay(1)
554
+
555
+ result = index - timedelta(1)
556
+ expected = index + timedelta(-1)
557
+ tm.assert_index_equal(result, expected)
558
+
559
+ def test_timedelta_tick_arithmetic(self):
560
+ # GH#4134, buggy with timedeltas
561
+ rng = pd.date_range("2013", "2014")
562
+ s = Series(rng)
563
+ result1 = rng - offsets.Hour(1)
564
+ result2 = DatetimeIndex(s - np.timedelta64(100000000, "ns"))
565
+ result3 = rng - np.timedelta64(100000000, "ns")
566
+ result4 = DatetimeIndex(s - offsets.Hour(1))
567
+
568
+ assert result1.freq == rng.freq
569
+ result1 = result1._with_freq(None)
570
+ tm.assert_index_equal(result1, result4)
571
+
572
+ assert result3.freq == rng.freq
573
+ result3 = result3._with_freq(None)
574
+ tm.assert_index_equal(result2, result3)
575
+
576
+ def test_tda_add_sub_index(self):
577
+ # Check that TimedeltaArray defers to Index on arithmetic ops
578
+ tdi = TimedeltaIndex(["1 days", NaT, "2 days"])
579
+ tda = tdi.array
580
+
581
+ dti = pd.date_range("1999-12-31", periods=3, freq="D")
582
+
583
+ result = tda + dti
584
+ expected = tdi + dti
585
+ tm.assert_index_equal(result, expected)
586
+
587
+ result = tda + tdi
588
+ expected = tdi + tdi
589
+ tm.assert_index_equal(result, expected)
590
+
591
+ result = tda - tdi
592
+ expected = tdi - tdi
593
+ tm.assert_index_equal(result, expected)
594
+
595
+ def test_tda_add_dt64_object_array(
596
+ self, performance_warning, box_with_array, tz_naive_fixture
597
+ ):
598
+ # Result should be cast back to DatetimeArray
599
+ box = box_with_array
600
+
601
+ dti = pd.date_range("2016-01-01", periods=3, tz=tz_naive_fixture)
602
+ dti = dti._with_freq(None)
603
+ tdi = dti - dti
604
+
605
+ obj = tm.box_expected(tdi, box)
606
+ other = tm.box_expected(dti, box)
607
+
608
+ with tm.assert_produces_warning(performance_warning):
609
+ result = obj + other.astype(object)
610
+ tm.assert_equal(result, other.astype(object))
611
+
612
+ # -------------------------------------------------------------
613
+ # Binary operations TimedeltaIndex and timedelta-like
614
+
615
+ def test_tdi_iadd_timedeltalike(self, two_hours, box_with_array):
616
+ # only test adding/sub offsets as + is now numeric
617
+ rng = timedelta_range("1 days", "10 days")
618
+ expected = timedelta_range("1 days 02:00:00", "10 days 02:00:00", freq="D")
619
+ if (
620
+ isinstance(two_hours, Timedelta)
621
+ and two_hours.unit == "ns"
622
+ and box_with_array is not pd.array
623
+ ):
624
+ # The EA op has to be _actually_ inplace so does not cast to a
625
+ # new dtype. For the others, the op can assign a new array
626
+ # and get the dtype that normally results from `rng + two_hours`
627
+ expected = expected.as_unit("ns")
628
+
629
+ rng = tm.box_expected(rng, box_with_array)
630
+ expected = tm.box_expected(expected, box_with_array)
631
+
632
+ orig_rng = rng
633
+ rng += two_hours
634
+ tm.assert_equal(rng, expected)
635
+ if box_with_array is not Index:
636
+ # Check that operation is actually inplace
637
+ tm.assert_equal(orig_rng, expected)
638
+
639
+ def test_tdi_isub_timedeltalike(self, two_hours, box_with_array):
640
+ # only test adding/sub offsets as - is now numeric
641
+ rng = timedelta_range("1 days", "10 days")
642
+ expected = timedelta_range("0 days 22:00:00", "9 days 22:00:00")
643
+ if (
644
+ isinstance(two_hours, Timedelta)
645
+ and two_hours.unit == "ns"
646
+ and box_with_array is not pd.array
647
+ ):
648
+ # The EA op has to be _actually_ inplace so does not cast to a
649
+ # new dtype. For the others, the op can assign a new array
650
+ # and get the dtype that normally results from `rng - two_hours`
651
+ expected = expected.as_unit("ns")
652
+
653
+ rng = tm.box_expected(rng, box_with_array)
654
+ expected = tm.box_expected(expected, box_with_array)
655
+
656
+ orig_rng = rng
657
+ rng -= two_hours
658
+ tm.assert_equal(rng, expected)
659
+ if box_with_array is not Index:
660
+ # Check that operation is actually inplace
661
+ tm.assert_equal(orig_rng, expected)
662
+
663
+ # -------------------------------------------------------------
664
+
665
+ def test_tdi_ops_attributes(self):
666
+ rng = timedelta_range("2 days", periods=5, freq="2D", name="x")
667
+
668
+ result = rng + 1 * rng.freq
669
+ exp = timedelta_range("4 days", periods=5, freq="2D", name="x")
670
+ tm.assert_index_equal(result, exp)
671
+ assert result.freq == "2D"
672
+
673
+ result = rng - 2 * rng.freq
674
+ exp = timedelta_range("-2 days", periods=5, freq="2D", name="x")
675
+ tm.assert_index_equal(result, exp)
676
+ assert result.freq == "2D"
677
+
678
+ result = rng * 2
679
+ exp = timedelta_range("4 days", periods=5, freq="4D", name="x")
680
+ tm.assert_index_equal(result, exp)
681
+ assert result.freq == "4D"
682
+
683
+ result = rng / 2
684
+ exp = timedelta_range("1 days", periods=5, freq="D", name="x")
685
+ tm.assert_index_equal(result, exp)
686
+ assert result.freq == "D"
687
+
688
+ result = -rng
689
+ exp = timedelta_range("-2 days", periods=5, freq="-2D", name="x")
690
+ tm.assert_index_equal(result, exp)
691
+ assert result.freq == "-2D"
692
+
693
+ rng = timedelta_range("-2 days", periods=5, freq="D", name="x")
694
+
695
+ result = abs(rng)
696
+ exp = TimedeltaIndex(
697
+ ["2 days", "1 days", "0 days", "1 days", "2 days"], name="x"
698
+ )
699
+ tm.assert_index_equal(result, exp)
700
+ assert result.freq is None
701
+
702
+
703
+ class TestAddSubNaTMasking:
704
+ # TODO: parametrize over boxes
705
+
706
+ @pytest.mark.parametrize("str_ts", ["1950-01-01", "1980-01-01"])
707
+ def test_tdarr_add_timestamp_nat_masking(self, box_with_array, str_ts):
708
+ # GH#17991 checking for overflow-masking with NaT
709
+ tdinat = pd.to_timedelta(["24658 days 11:15:00", "NaT"])
710
+ tdobj = tm.box_expected(tdinat, box_with_array)
711
+
712
+ ts = Timestamp(str_ts)
713
+ ts_variants = [
714
+ ts,
715
+ ts.to_pydatetime(),
716
+ ts.to_datetime64().astype("datetime64[ns]"),
717
+ ts.to_datetime64().astype("datetime64[D]"),
718
+ ]
719
+
720
+ for variant in ts_variants:
721
+ res = tdobj + variant
722
+ if box_with_array is DataFrame:
723
+ assert res.iloc[1, 1] is NaT
724
+ else:
725
+ assert res[1] is NaT
726
+
727
+ def test_tdi_add_overflow(self):
728
+ # See GH#14068
729
+ # preliminary test scalar analogue of vectorized tests below
730
+ # TODO: Make raised error message more informative and test
731
+ ts = Timestamp("2000").as_unit("ns")
732
+ with pytest.raises(OutOfBoundsDatetime, match="10155196800000000000"):
733
+ pd.to_timedelta(106580, "D") + ts
734
+ with pytest.raises(OutOfBoundsDatetime, match="10155196800000000000"):
735
+ ts + pd.to_timedelta(106580, "D")
736
+
737
+ _NaT = NaT._value + 1
738
+ td = pd.to_timedelta([106580], "D").as_unit("ns")
739
+ msg = "Overflow in int64 addition"
740
+ with pytest.raises(OverflowError, match=msg):
741
+ td + Timestamp("2000")
742
+ with pytest.raises(OverflowError, match=msg):
743
+ Timestamp("2000") + td
744
+ with pytest.raises(OverflowError, match=msg):
745
+ pd.to_timedelta([_NaT]) - Timedelta("1 days")
746
+ with pytest.raises(OverflowError, match=msg):
747
+ pd.to_timedelta(["5 days", _NaT]) - Timedelta("1 days")
748
+ with pytest.raises(OverflowError, match=msg):
749
+ (
750
+ pd.to_timedelta([_NaT, "5 days", "1 hours"])
751
+ - pd.to_timedelta(["7 seconds", _NaT, "4 hours"])
752
+ )
753
+
754
+ # These should not overflow!
755
+ exp = TimedeltaIndex([NaT], dtype="m8[us]")
756
+ result = pd.to_timedelta([NaT]) - Timedelta("1 days")
757
+ tm.assert_index_equal(result, exp)
758
+
759
+ exp = TimedeltaIndex(["4 days", NaT])
760
+ result = pd.to_timedelta(["5 days", NaT]) - Timedelta("1 days")
761
+ tm.assert_index_equal(result, exp)
762
+
763
+ exp = TimedeltaIndex([NaT, NaT, "5 hours"])
764
+ result = pd.to_timedelta([NaT, "5 days", "1 hours"]) + pd.to_timedelta(
765
+ ["7 seconds", NaT, "4 hours"]
766
+ )
767
+ tm.assert_index_equal(result, exp)
768
+
769
+
770
+ class TestTimedeltaArraylikeAddSubOps:
771
+ # Tests for timedelta64[ns] __add__, __sub__, __radd__, __rsub__
772
+
773
+ def test_sub_nat_retain_unit(self):
774
+ ser = pd.to_timedelta(Series(["00:00:01"])).astype("m8[s]")
775
+
776
+ result = ser - NaT
777
+ expected = Series([NaT], dtype="m8[s]")
778
+ tm.assert_series_equal(result, expected)
779
+
780
+ # TODO: moved from tests.indexes.timedeltas.test_arithmetic; needs
781
+ # parametrization+de-duplication
782
+ def test_timedelta_ops_with_missing_values(self):
783
+ # setup
784
+ s1 = pd.to_timedelta(Series(["00:00:01"]))
785
+ s2 = pd.to_timedelta(Series(["00:00:02"]))
786
+
787
+ sn = pd.to_timedelta(Series([NaT], dtype="m8[us]"))
788
+
789
+ df1 = DataFrame(["00:00:01"]).apply(pd.to_timedelta)
790
+ df2 = DataFrame(["00:00:02"]).apply(pd.to_timedelta)
791
+
792
+ dfn = DataFrame([NaT._value]).apply(pd.to_timedelta).astype("m8[us]")
793
+
794
+ scalar1 = pd.to_timedelta("00:00:01")
795
+ scalar2 = pd.to_timedelta("00:00:02")
796
+ timedelta_NaT = pd.to_timedelta("NaT")
797
+
798
+ actual = scalar1 + scalar1
799
+ assert actual == scalar2
800
+ actual = scalar2 - scalar1
801
+ assert actual == scalar1
802
+
803
+ actual = s1 + s1
804
+ tm.assert_series_equal(actual, s2)
805
+ actual = s2 - s1
806
+ tm.assert_series_equal(actual, s1)
807
+
808
+ actual = s1 + scalar1
809
+ tm.assert_series_equal(actual, s2)
810
+ actual = scalar1 + s1
811
+ tm.assert_series_equal(actual, s2)
812
+ actual = s2 - scalar1
813
+ tm.assert_series_equal(actual, s1)
814
+ actual = -scalar1 + s2
815
+ tm.assert_series_equal(actual, s1)
816
+
817
+ actual = s1 + timedelta_NaT
818
+ tm.assert_series_equal(actual, sn)
819
+ actual = timedelta_NaT + s1
820
+ tm.assert_series_equal(actual, sn)
821
+ actual = s1 - timedelta_NaT
822
+ tm.assert_series_equal(actual, sn)
823
+ actual = -timedelta_NaT + s1
824
+ tm.assert_series_equal(actual, sn)
825
+
826
+ msg = "unsupported operand type"
827
+ with pytest.raises(TypeError, match=msg):
828
+ s1 + np.nan
829
+ with pytest.raises(TypeError, match=msg):
830
+ np.nan + s1
831
+ with pytest.raises(TypeError, match=msg):
832
+ s1 - np.nan
833
+ with pytest.raises(TypeError, match=msg):
834
+ -np.nan + s1
835
+
836
+ actual = s1 + NaT
837
+ tm.assert_series_equal(actual, sn)
838
+ actual = s2 - NaT
839
+ tm.assert_series_equal(actual, sn)
840
+
841
+ actual = s1 + df1
842
+ tm.assert_frame_equal(actual, df2)
843
+ actual = s2 - df1
844
+ tm.assert_frame_equal(actual, df1)
845
+ actual = df1 + s1
846
+ tm.assert_frame_equal(actual, df2)
847
+ actual = df2 - s1
848
+ tm.assert_frame_equal(actual, df1)
849
+
850
+ actual = df1 + df1
851
+ tm.assert_frame_equal(actual, df2)
852
+ actual = df2 - df1
853
+ tm.assert_frame_equal(actual, df1)
854
+
855
+ actual = df1 + scalar1
856
+ tm.assert_frame_equal(actual, df2)
857
+ actual = df2 - scalar1
858
+ tm.assert_frame_equal(actual, df1)
859
+
860
+ actual = df1 + timedelta_NaT
861
+ tm.assert_frame_equal(actual, dfn)
862
+ actual = df1 - timedelta_NaT
863
+ tm.assert_frame_equal(actual, dfn)
864
+
865
+ msg = "cannot subtract a datelike from|unsupported operand type"
866
+ with pytest.raises(TypeError, match=msg):
867
+ df1 + np.nan
868
+ with pytest.raises(TypeError, match=msg):
869
+ df1 - np.nan
870
+
871
+ actual = df1 + NaT # NaT is datetime, not timedelta
872
+ tm.assert_frame_equal(actual, dfn)
873
+ actual = df1 - NaT
874
+ tm.assert_frame_equal(actual, dfn)
875
+
876
+ # TODO: moved from tests.series.test_operators, needs splitting, cleanup,
877
+ # de-duplication, box-parametrization...
878
+ def test_operators_timedelta64(self):
879
+ # series ops
880
+ v1 = pd.date_range("2012-1-1", periods=3, freq="D", unit="ns")
881
+ v2 = pd.date_range("2012-1-2", periods=3, freq="D", unit="ns")
882
+ rs = Series(v2) - Series(v1)
883
+ xp = Series(1e9 * 3600 * 24, rs.index).astype("int64").astype("timedelta64[ns]")
884
+ tm.assert_series_equal(rs, xp)
885
+ assert rs.dtype == "timedelta64[ns]"
886
+
887
+ df = DataFrame({"A": v1})
888
+ td = Series([timedelta(days=i) for i in range(3)], dtype="m8[ns]")
889
+ assert td.dtype == "timedelta64[ns]"
890
+
891
+ # series on the rhs
892
+ result = df["A"] - df["A"].shift()
893
+ assert result.dtype == "timedelta64[ns]"
894
+
895
+ result = df["A"] + td
896
+ assert result.dtype == "M8[ns]"
897
+
898
+ # scalar Timestamp on rhs
899
+ maxa = df["A"].max()
900
+ assert isinstance(maxa, Timestamp)
901
+
902
+ resultb = df["A"] - df["A"].max()
903
+ assert resultb.dtype == "timedelta64[ns]"
904
+
905
+ # timestamp on lhs
906
+ result = resultb + df["A"]
907
+ values = [Timestamp("20111230"), Timestamp("20120101"), Timestamp("20120103")]
908
+ expected = Series(values, dtype="M8[ns]", name="A")
909
+ tm.assert_series_equal(result, expected)
910
+
911
+ # datetimes on rhs
912
+ result = df["A"] - datetime(2001, 1, 1)
913
+ expected = Series(
914
+ [timedelta(days=4017 + i) for i in range(3)], name="A", dtype="m8[ns]"
915
+ )
916
+ tm.assert_series_equal(result, expected)
917
+ assert result.dtype == "m8[ns]"
918
+
919
+ d = datetime(2001, 1, 1, 3, 4)
920
+ resulta = df["A"] - d
921
+ assert resulta.dtype == "m8[ns]"
922
+
923
+ # roundtrip
924
+ resultb = resulta + d
925
+ tm.assert_series_equal(df["A"], resultb)
926
+
927
+ # timedeltas on rhs
928
+ td = timedelta(days=1)
929
+ resulta = df["A"] + td
930
+ resultb = resulta - td
931
+ tm.assert_series_equal(resultb, df["A"])
932
+ assert resultb.dtype == "M8[ns]"
933
+
934
+ # roundtrip
935
+ td = timedelta(minutes=5, seconds=3)
936
+ resulta = df["A"] + td
937
+ resultb = resulta - td
938
+ tm.assert_series_equal(df["A"], resultb)
939
+ assert resultb.dtype == "M8[ns]"
940
+
941
+ # inplace
942
+ value = rs[2] + np.timedelta64(timedelta(minutes=5, seconds=1))
943
+ rs[2] += np.timedelta64(timedelta(minutes=5, seconds=1))
944
+ assert rs[2] == value
945
+
946
+ def test_timedelta64_ops_nat(self):
947
+ # GH 11349
948
+ timedelta_series = Series([NaT, Timedelta("1s")])
949
+ nat_series_dtype_timedelta = Series([NaT, NaT], dtype="timedelta64[us]")
950
+ single_nat_dtype_timedelta = Series([NaT], dtype="timedelta64[us]")
951
+
952
+ # subtraction
953
+ tm.assert_series_equal(timedelta_series - NaT, nat_series_dtype_timedelta)
954
+ tm.assert_series_equal(-NaT + timedelta_series, nat_series_dtype_timedelta)
955
+
956
+ tm.assert_series_equal(
957
+ timedelta_series - single_nat_dtype_timedelta, nat_series_dtype_timedelta
958
+ )
959
+ tm.assert_series_equal(
960
+ -single_nat_dtype_timedelta + timedelta_series, nat_series_dtype_timedelta
961
+ )
962
+
963
+ # addition
964
+ tm.assert_series_equal(
965
+ nat_series_dtype_timedelta + NaT, nat_series_dtype_timedelta
966
+ )
967
+ tm.assert_series_equal(
968
+ NaT + nat_series_dtype_timedelta, nat_series_dtype_timedelta
969
+ )
970
+
971
+ tm.assert_series_equal(
972
+ nat_series_dtype_timedelta + single_nat_dtype_timedelta,
973
+ nat_series_dtype_timedelta,
974
+ )
975
+ tm.assert_series_equal(
976
+ single_nat_dtype_timedelta + nat_series_dtype_timedelta,
977
+ nat_series_dtype_timedelta,
978
+ )
979
+
980
+ tm.assert_series_equal(timedelta_series + NaT, nat_series_dtype_timedelta)
981
+ tm.assert_series_equal(NaT + timedelta_series, nat_series_dtype_timedelta)
982
+
983
+ tm.assert_series_equal(
984
+ timedelta_series + single_nat_dtype_timedelta, nat_series_dtype_timedelta
985
+ )
986
+ tm.assert_series_equal(
987
+ single_nat_dtype_timedelta + timedelta_series, nat_series_dtype_timedelta
988
+ )
989
+
990
+ tm.assert_series_equal(
991
+ nat_series_dtype_timedelta + NaT, nat_series_dtype_timedelta
992
+ )
993
+ tm.assert_series_equal(
994
+ NaT + nat_series_dtype_timedelta, nat_series_dtype_timedelta
995
+ )
996
+
997
+ tm.assert_series_equal(
998
+ nat_series_dtype_timedelta + single_nat_dtype_timedelta,
999
+ nat_series_dtype_timedelta,
1000
+ )
1001
+ tm.assert_series_equal(
1002
+ single_nat_dtype_timedelta + nat_series_dtype_timedelta,
1003
+ nat_series_dtype_timedelta,
1004
+ )
1005
+
1006
+ # multiplication
1007
+ tm.assert_series_equal(
1008
+ nat_series_dtype_timedelta * 1.0, nat_series_dtype_timedelta
1009
+ )
1010
+ tm.assert_series_equal(
1011
+ 1.0 * nat_series_dtype_timedelta, nat_series_dtype_timedelta
1012
+ )
1013
+
1014
+ tm.assert_series_equal(timedelta_series * 1, timedelta_series)
1015
+ tm.assert_series_equal(1 * timedelta_series, timedelta_series)
1016
+
1017
+ tm.assert_series_equal(timedelta_series * 1.5, Series([NaT, Timedelta("1.5s")]))
1018
+ tm.assert_series_equal(1.5 * timedelta_series, Series([NaT, Timedelta("1.5s")]))
1019
+
1020
+ tm.assert_series_equal(timedelta_series * np.nan, nat_series_dtype_timedelta)
1021
+ tm.assert_series_equal(np.nan * timedelta_series, nat_series_dtype_timedelta)
1022
+
1023
+ # division
1024
+ tm.assert_series_equal(timedelta_series / 2, Series([NaT, Timedelta("0.5s")]))
1025
+ tm.assert_series_equal(timedelta_series / 2.0, Series([NaT, Timedelta("0.5s")]))
1026
+ tm.assert_series_equal(timedelta_series / np.nan, nat_series_dtype_timedelta)
1027
+
1028
+ # -------------------------------------------------------------
1029
+ # Binary operations td64 arraylike and datetime-like
1030
+
1031
+ @pytest.mark.parametrize("cls", [Timestamp, datetime, np.datetime64])
1032
+ def test_td64arr_add_sub_datetimelike_scalar(
1033
+ self, cls, box_with_array, tz_naive_fixture
1034
+ ):
1035
+ # GH#11925, GH#29558, GH#23215
1036
+ tz = tz_naive_fixture
1037
+
1038
+ dt_scalar = Timestamp("2012-01-01", tz=tz)
1039
+ if cls is datetime:
1040
+ ts = dt_scalar.to_pydatetime()
1041
+ elif cls is np.datetime64:
1042
+ if tz_naive_fixture is not None:
1043
+ pytest.skip(f"{cls} doesn't support {tz_naive_fixture}")
1044
+ ts = dt_scalar.to_datetime64()
1045
+ else:
1046
+ ts = dt_scalar
1047
+
1048
+ tdi = timedelta_range("1 day", periods=3)
1049
+ expected = pd.date_range("2012-01-02", periods=3, tz=tz)
1050
+ if tz is not None and not timezones.is_utc(expected.tz):
1051
+ # Day is no longer preserved by timedelta add/sub in pandas3 because
1052
+ # it represents Calendar-Day instead of 24h
1053
+ expected = expected._with_freq(None)
1054
+
1055
+ tdarr = tm.box_expected(tdi, box_with_array)
1056
+ expected = tm.box_expected(expected, box_with_array)
1057
+
1058
+ tm.assert_equal(ts + tdarr, expected)
1059
+ tm.assert_equal(tdarr + ts, expected)
1060
+
1061
+ expected2 = pd.date_range("2011-12-31", periods=3, freq="-1D", tz=tz)
1062
+ if tz is not None and not timezones.is_utc(expected2.tz):
1063
+ # Day is no longer preserved by timedelta add/sub in pandas3 because
1064
+ # it represents Calendar-Day instead of 24h
1065
+ expected2 = expected2._with_freq(None)
1066
+ expected2 = tm.box_expected(expected2, box_with_array)
1067
+
1068
+ tm.assert_equal(ts - tdarr, expected2)
1069
+ tm.assert_equal(ts + (-tdarr), expected2)
1070
+
1071
+ msg = "cannot subtract a datelike"
1072
+ with pytest.raises(TypeError, match=msg):
1073
+ tdarr - ts
1074
+
1075
+ @pytest.mark.filterwarnings("ignore:.*generic.*unit.*:DeprecationWarning")
1076
+ def test_td64arr_add_datetime64_nat(self, box_with_array):
1077
+ # GH#23215
1078
+ # The bare np.datetime64("NaT") is load-bearing here: the expected
1079
+ # M8[us] dtype comes from numpy's default unit promotion from the
1080
+ # generic-unit scalar. NumPy nightly emits a DeprecationWarning for
1081
+ # the bare form, which the filterwarnings marker ignores.
1082
+ other = np.datetime64("NaT")
1083
+
1084
+ tdi = timedelta_range("1 day", periods=3)
1085
+ expected = DatetimeIndex(["NaT", "NaT", "NaT"], dtype="M8[us]")
1086
+
1087
+ tdser = tm.box_expected(tdi, box_with_array)
1088
+ expected = tm.box_expected(expected, box_with_array)
1089
+
1090
+ tm.assert_equal(tdser + other, expected)
1091
+ tm.assert_equal(other + tdser, expected)
1092
+
1093
+ def test_td64arr_sub_dt64_array(self, box_with_array):
1094
+ dti = pd.date_range("2016-01-01", periods=3)
1095
+ tdi = TimedeltaIndex(["-1 Day"] * 3)
1096
+ dtarr = dti.values
1097
+ expected = DatetimeIndex(dtarr) - tdi
1098
+
1099
+ tdi = tm.box_expected(tdi, box_with_array)
1100
+ expected = tm.box_expected(expected, box_with_array)
1101
+
1102
+ msg = "cannot subtract a datelike from"
1103
+ with pytest.raises(TypeError, match=msg):
1104
+ tdi - dtarr
1105
+
1106
+ # TimedeltaIndex.__rsub__
1107
+ result = dtarr - tdi
1108
+ tm.assert_equal(result, expected)
1109
+
1110
+ def test_td64arr_add_dt64_array(self, box_with_array):
1111
+ dti = pd.date_range("2016-01-01", periods=3)
1112
+ tdi = TimedeltaIndex(["-1 Day"] * 3)
1113
+ dtarr = dti.values
1114
+ expected = DatetimeIndex(dtarr) + tdi
1115
+
1116
+ tdi = tm.box_expected(tdi, box_with_array)
1117
+ expected = tm.box_expected(expected, box_with_array)
1118
+
1119
+ result = tdi + dtarr
1120
+ tm.assert_equal(result, expected)
1121
+ result = dtarr + tdi
1122
+ tm.assert_equal(result, expected)
1123
+
1124
+ # ------------------------------------------------------------------
1125
+ # Invalid __add__/__sub__ operations
1126
+
1127
+ @pytest.mark.parametrize("pi_freq", ["D", "W", "Q", "h"])
1128
+ @pytest.mark.parametrize("tdi_freq", [None, "h"])
1129
+ def test_td64arr_sub_periodlike(
1130
+ self, box_with_array, box_with_array2, tdi_freq, pi_freq
1131
+ ):
1132
+ # GH#20049 subtracting PeriodIndex should raise TypeError
1133
+ tdi = TimedeltaIndex(["1 hours", "2 hours"], freq=tdi_freq)
1134
+ dti = Timestamp("2018-03-07 17:16:40") + tdi
1135
+ pi = dti.to_period(pi_freq)
1136
+ per = pi[0]
1137
+
1138
+ tdi = tm.box_expected(tdi, box_with_array)
1139
+ pi = tm.box_expected(pi, box_with_array2)
1140
+ msg = "|".join(
1141
+ [
1142
+ "cannot subtract",
1143
+ "unsupported operand type",
1144
+ r"bad operand type for unary \-: 'PeriodArray'",
1145
+ r"Input has different freq=-1h from PeriodArray\(.*\)",
1146
+ "Cannot add/subtract timedelta-like from PeriodArray",
1147
+ ]
1148
+ )
1149
+ with pytest.raises(TypeError, match=msg):
1150
+ tdi - pi
1151
+
1152
+ # GH#13078 subtraction of Period scalar not supported
1153
+ with pytest.raises(TypeError, match=msg):
1154
+ tdi - per
1155
+
1156
+ @pytest.mark.parametrize(
1157
+ "other",
1158
+ [
1159
+ # GH#12624 for str case
1160
+ "a",
1161
+ # GH#19123
1162
+ 1,
1163
+ 1.5,
1164
+ np.array(2),
1165
+ ],
1166
+ )
1167
+ def test_td64arr_addsub_numeric_scalar_invalid(self, box_with_array, other):
1168
+ # vector-like others are tested in test_td64arr_add_sub_numeric_arr_invalid
1169
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
1170
+ tdarr = tm.box_expected(tdser, box_with_array)
1171
+
1172
+ assert_invalid_addsub_type(tdarr, other)
1173
+
1174
+ @pytest.mark.parametrize(
1175
+ "vec",
1176
+ [
1177
+ np.array([1, 2, 3]),
1178
+ Index([1, 2, 3]),
1179
+ Series([1, 2, 3]),
1180
+ DataFrame([[1, 2, 3]]),
1181
+ ],
1182
+ ids=lambda x: type(x).__name__,
1183
+ )
1184
+ def test_td64arr_addsub_numeric_arr_invalid(
1185
+ self, box_with_array, vec, any_real_numpy_dtype
1186
+ ):
1187
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
1188
+ tdarr = tm.box_expected(tdser, box_with_array)
1189
+
1190
+ vector = vec.astype(any_real_numpy_dtype)
1191
+ assert_invalid_addsub_type(tdarr, vector)
1192
+
1193
+ def test_td64arr_add_sub_int(self, box_with_array, one):
1194
+ # Variants of `one` for #19012, deprecated GH#22535
1195
+ rng = timedelta_range("1 days 09:00:00", freq="h", periods=10)
1196
+ tdarr = tm.box_expected(rng, box_with_array)
1197
+
1198
+ msg = "Addition/subtraction of integers"
1199
+ assert_invalid_addsub_type(tdarr, one, msg)
1200
+
1201
+ # TODO: get inplace ops into assert_invalid_addsub_type
1202
+ with pytest.raises(TypeError, match=msg):
1203
+ tdarr += one
1204
+ with pytest.raises(TypeError, match=msg):
1205
+ tdarr -= one
1206
+
1207
+ def test_td64arr_add_sub_integer_array(self, box_with_array):
1208
+ # GH#19959, deprecated GH#22535
1209
+ # GH#22696 for DataFrame case, check that we don't dispatch to numpy
1210
+ # implementation, which treats int64 as m8[ns]
1211
+ box = box_with_array
1212
+ xbox = np.ndarray if box is pd.array else box
1213
+
1214
+ rng = timedelta_range("1 days 09:00:00", freq="h", periods=3)
1215
+ tdarr = tm.box_expected(rng, box)
1216
+ other = tm.box_expected([4, 3, 2], xbox)
1217
+
1218
+ msg = "Addition/subtraction of integers and integer-arrays"
1219
+ assert_invalid_addsub_type(tdarr, other, msg)
1220
+
1221
+ def test_td64arr_addsub_integer_array_no_freq(self, box_with_array):
1222
+ # GH#19959
1223
+ box = box_with_array
1224
+ xbox = np.ndarray if box is pd.array else box
1225
+
1226
+ tdi = TimedeltaIndex(["1 Day", "NaT", "3 Hours"])
1227
+ tdarr = tm.box_expected(tdi, box)
1228
+ other = tm.box_expected([14, -1, 16], xbox)
1229
+
1230
+ msg = "Addition/subtraction of integers"
1231
+ assert_invalid_addsub_type(tdarr, other, msg)
1232
+
1233
+ # ------------------------------------------------------------------
1234
+ # Operations with timedelta-like others
1235
+
1236
+ def test_td64arr_add_sub_td64_array(self, box_with_array):
1237
+ box = box_with_array
1238
+ dti = pd.date_range("2016-01-01", periods=3)
1239
+ tdi = dti - dti.shift(1)
1240
+ tdarr = tdi.values
1241
+
1242
+ expected = 2 * tdi
1243
+ tdi = tm.box_expected(tdi, box)
1244
+ expected = tm.box_expected(expected, box)
1245
+
1246
+ result = tdi + tdarr
1247
+ tm.assert_equal(result, expected)
1248
+ result = tdarr + tdi
1249
+ tm.assert_equal(result, expected)
1250
+
1251
+ expected_sub = 0 * tdi
1252
+ result = tdi - tdarr
1253
+ tm.assert_equal(result, expected_sub)
1254
+ result = tdarr - tdi
1255
+ tm.assert_equal(result, expected_sub)
1256
+
1257
+ def test_td64arr_add_sub_tdi(self, box_with_array, names):
1258
+ # GH#17250 make sure result dtype is correct
1259
+ # GH#19043 make sure names are propagated correctly
1260
+ box = box_with_array
1261
+ exname = get_expected_name(box, names)
1262
+
1263
+ tdi = TimedeltaIndex(["0 days", "1 day"], name=names[1])
1264
+ tdi = np.array(tdi) if box in [tm.to_array, pd.array] else tdi
1265
+ ser = Series(
1266
+ [Timedelta(hours=3), Timedelta(hours=4)], name=names[0], dtype="m8[ns]"
1267
+ )
1268
+ expected = Series(
1269
+ [Timedelta(hours=3), Timedelta(days=1, hours=4)],
1270
+ name=exname,
1271
+ dtype="m8[ns]",
1272
+ )
1273
+
1274
+ ser = tm.box_expected(ser, box)
1275
+ expected = tm.box_expected(expected, box)
1276
+
1277
+ result = tdi + ser
1278
+ tm.assert_equal(result, expected)
1279
+ assert_dtype(result, "timedelta64[ns]")
1280
+
1281
+ result = ser + tdi
1282
+ tm.assert_equal(result, expected)
1283
+ assert_dtype(result, "timedelta64[ns]")
1284
+
1285
+ expected = Series(
1286
+ [Timedelta(hours=-3), Timedelta(days=1, hours=-4)],
1287
+ name=exname,
1288
+ dtype="m8[ns]",
1289
+ )
1290
+ expected = tm.box_expected(expected, box)
1291
+
1292
+ result = tdi - ser
1293
+ tm.assert_equal(result, expected)
1294
+ assert_dtype(result, "timedelta64[ns]")
1295
+
1296
+ result = ser - tdi
1297
+ tm.assert_equal(result, -expected)
1298
+ assert_dtype(result, "timedelta64[ns]")
1299
+
1300
+ @pytest.mark.parametrize("tdnat", [np.timedelta64("NaT", "us"), NaT])
1301
+ def test_td64arr_add_sub_td64_nat(self, box_with_array, tdnat):
1302
+ # GH#18808, GH#23320 special handling for timedelta64("NaT")
1303
+ box = box_with_array
1304
+ tdi = TimedeltaIndex([NaT, Timedelta("1s")])
1305
+ expected = TimedeltaIndex(["NaT"] * 2).as_unit("us")
1306
+
1307
+ obj = tm.box_expected(tdi, box)
1308
+ expected = tm.box_expected(expected, box)
1309
+
1310
+ result = obj + tdnat
1311
+ tm.assert_equal(result, expected)
1312
+ result = tdnat + obj
1313
+ tm.assert_equal(result, expected)
1314
+ result = obj - tdnat
1315
+ tm.assert_equal(result, expected)
1316
+ result = tdnat - obj
1317
+ tm.assert_equal(result, expected)
1318
+
1319
+ def test_td64arr_add_timedeltalike(self, two_hours, box_with_array):
1320
+ # only test adding/sub offsets as + is now numeric
1321
+ # GH#10699 for Tick cases
1322
+ box = box_with_array
1323
+ rng = timedelta_range("1 days", "10 days")
1324
+ expected = timedelta_range("1 days 02:00:00", "10 days 02:00:00", freq="D")
1325
+ if isinstance(two_hours, Timedelta) and two_hours.unit == "ns":
1326
+ expected = expected.as_unit("ns")
1327
+
1328
+ rng = tm.box_expected(rng, box)
1329
+ expected = tm.box_expected(expected, box)
1330
+
1331
+ result = rng + two_hours
1332
+ tm.assert_equal(result, expected)
1333
+
1334
+ result = two_hours + rng
1335
+ tm.assert_equal(result, expected)
1336
+
1337
+ def test_td64arr_sub_timedeltalike(self, two_hours, box_with_array):
1338
+ # only test adding/sub offsets as - is now numeric
1339
+ # GH#10699 for Tick cases
1340
+ box = box_with_array
1341
+ rng = timedelta_range("1 days", "10 days")
1342
+ expected = timedelta_range("0 days 22:00:00", "9 days 22:00:00")
1343
+ if isinstance(two_hours, Timedelta) and two_hours.unit == "ns":
1344
+ expected = expected.as_unit("ns")
1345
+
1346
+ rng = tm.box_expected(rng, box)
1347
+ expected = tm.box_expected(expected, box)
1348
+
1349
+ result = rng - two_hours
1350
+ tm.assert_equal(result, expected)
1351
+
1352
+ result = two_hours - rng
1353
+ tm.assert_equal(result, -expected)
1354
+
1355
+ # ------------------------------------------------------------------
1356
+ # __add__/__sub__ with DateOffsets and arrays of DateOffsets
1357
+
1358
+ def test_td64arr_add_sub_offset_index(
1359
+ self, performance_warning, names, box_with_array
1360
+ ):
1361
+ # GH#18849, GH#19744
1362
+ box = box_with_array
1363
+ exname = get_expected_name(box, names)
1364
+
1365
+ tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"], name=names[0])
1366
+ other = Index([offsets.Hour(n=1), offsets.Minute(n=-2)], name=names[1])
1367
+ other = np.array(other) if box in [tm.to_array, pd.array] else other
1368
+
1369
+ expected = TimedeltaIndex(
1370
+ [tdi[n] + other[n] for n in range(len(tdi))], freq="infer", name=exname
1371
+ )
1372
+ expected_sub = TimedeltaIndex(
1373
+ [tdi[n] - other[n] for n in range(len(tdi))], freq="infer", name=exname
1374
+ )
1375
+
1376
+ tdi = tm.box_expected(tdi, box)
1377
+ expected = tm.box_expected(expected, box).astype(object)
1378
+ expected_sub = tm.box_expected(expected_sub, box).astype(object)
1379
+
1380
+ with tm.assert_produces_warning(performance_warning):
1381
+ res = tdi + other
1382
+ tm.assert_equal(res, expected)
1383
+
1384
+ with tm.assert_produces_warning(performance_warning):
1385
+ res2 = other + tdi
1386
+ tm.assert_equal(res2, expected)
1387
+
1388
+ with tm.assert_produces_warning(performance_warning):
1389
+ res_sub = tdi - other
1390
+ tm.assert_equal(res_sub, expected_sub)
1391
+
1392
+ def test_td64arr_add_sub_offset_array(self, performance_warning, box_with_array):
1393
+ # GH#18849, GH#18824
1394
+ box = box_with_array
1395
+ tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"])
1396
+ other = np.array([offsets.Hour(n=1), offsets.Minute(n=-2)])
1397
+
1398
+ expected = TimedeltaIndex(
1399
+ [tdi[n] + other[n] for n in range(len(tdi))], freq="infer"
1400
+ )
1401
+ expected_sub = TimedeltaIndex(
1402
+ [tdi[n] - other[n] for n in range(len(tdi))], freq="infer"
1403
+ )
1404
+
1405
+ tdi = tm.box_expected(tdi, box)
1406
+ expected = tm.box_expected(expected, box).astype(object)
1407
+
1408
+ with tm.assert_produces_warning(performance_warning):
1409
+ res = tdi + other
1410
+ tm.assert_equal(res, expected)
1411
+
1412
+ with tm.assert_produces_warning(performance_warning):
1413
+ res2 = other + tdi
1414
+ tm.assert_equal(res2, expected)
1415
+
1416
+ expected_sub = tm.box_expected(expected_sub, box_with_array).astype(object)
1417
+ with tm.assert_produces_warning(performance_warning):
1418
+ res_sub = tdi - other
1419
+ tm.assert_equal(res_sub, expected_sub)
1420
+
1421
+ def test_td64arr_with_offset_series(
1422
+ self, performance_warning, names, box_with_array
1423
+ ):
1424
+ # GH#18849
1425
+ box = box_with_array
1426
+ box2 = Series if box in [Index, tm.to_array, pd.array] else box
1427
+ exname = get_expected_name(box, names)
1428
+
1429
+ tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"], name=names[0])
1430
+ other = Series([offsets.Hour(n=1), offsets.Minute(n=-2)], name=names[1])
1431
+
1432
+ expected_add = Series(
1433
+ [tdi[n] + other[n] for n in range(len(tdi))], name=exname, dtype=object
1434
+ )
1435
+ obj = tm.box_expected(tdi, box)
1436
+ expected_add = tm.box_expected(expected_add, box2).astype(object)
1437
+
1438
+ with tm.assert_produces_warning(performance_warning):
1439
+ res = obj + other
1440
+ tm.assert_equal(res, expected_add)
1441
+
1442
+ with tm.assert_produces_warning(performance_warning):
1443
+ res2 = other + obj
1444
+ tm.assert_equal(res2, expected_add)
1445
+
1446
+ expected_sub = Series(
1447
+ [tdi[n] - other[n] for n in range(len(tdi))], name=exname, dtype=object
1448
+ )
1449
+ expected_sub = tm.box_expected(expected_sub, box2).astype(object)
1450
+
1451
+ with tm.assert_produces_warning(performance_warning):
1452
+ res3 = obj - other
1453
+ tm.assert_equal(res3, expected_sub)
1454
+
1455
+ @pytest.mark.parametrize("obox", [np.array, Index, Series])
1456
+ def test_td64arr_addsub_anchored_offset_arraylike(
1457
+ self, performance_warning, obox, box_with_array
1458
+ ):
1459
+ # GH#18824
1460
+ tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"])
1461
+ tdi = tm.box_expected(tdi, box_with_array)
1462
+
1463
+ anchored = obox([offsets.MonthEnd(), offsets.Day(n=2)])
1464
+
1465
+ # addition/subtraction ops with anchored offsets should issue
1466
+ # a PerformanceWarning and _then_ raise a TypeError.
1467
+ msg = "has incorrect type|cannot add the type MonthEnd"
1468
+ with pytest.raises(TypeError, match=msg):
1469
+ with tm.assert_produces_warning(performance_warning):
1470
+ tdi + anchored
1471
+ with pytest.raises(TypeError, match=msg):
1472
+ with tm.assert_produces_warning(performance_warning):
1473
+ anchored + tdi
1474
+ with pytest.raises(TypeError, match=msg):
1475
+ with tm.assert_produces_warning(performance_warning):
1476
+ tdi - anchored
1477
+ with pytest.raises(TypeError, match=msg):
1478
+ with tm.assert_produces_warning(performance_warning):
1479
+ anchored - tdi
1480
+
1481
+ # ------------------------------------------------------------------
1482
+ # Unsorted
1483
+
1484
+ def test_td64arr_add_sub_object_array(self, performance_warning, box_with_array):
1485
+ box = box_with_array
1486
+ xbox = np.ndarray if box is pd.array else box
1487
+
1488
+ tdi = timedelta_range("1 day", periods=3, freq="D")
1489
+ tdarr = tm.box_expected(tdi, box)
1490
+
1491
+ other = np.array([Timedelta(days=1), offsets.Day(2), Timestamp("2000-01-04")])
1492
+
1493
+ with tm.assert_produces_warning(performance_warning):
1494
+ result = tdarr + other
1495
+
1496
+ expected = Index(
1497
+ [Timedelta(days=2), Timedelta(days=4), Timestamp("2000-01-07")]
1498
+ )
1499
+ expected = tm.box_expected(expected, xbox).astype(object)
1500
+ tm.assert_equal(result, expected)
1501
+
1502
+ msg = "unsupported operand type|cannot subtract a datelike"
1503
+ with pytest.raises(TypeError, match=msg):
1504
+ with tm.assert_produces_warning(performance_warning):
1505
+ tdarr - other
1506
+
1507
+ with tm.assert_produces_warning(performance_warning):
1508
+ result = other - tdarr
1509
+
1510
+ expected = Index([Timedelta(0), Timedelta(0), Timestamp("2000-01-01")])
1511
+ expected = tm.box_expected(expected, xbox).astype(object)
1512
+ tm.assert_equal(result, expected)
1513
+
1514
+
1515
+ class TestTimedeltaArraylikeMulDivOps:
1516
+ # Tests for timedelta64[ns]
1517
+ # __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__
1518
+
1519
+ # ------------------------------------------------------------------
1520
+ # Multiplication
1521
+ # organized with scalar others first, then array-like
1522
+
1523
+ def test_td64arr_mul_int(self, box_with_array):
1524
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
1525
+ idx = tm.box_expected(idx, box_with_array)
1526
+
1527
+ result = idx * 1
1528
+ tm.assert_equal(result, idx)
1529
+
1530
+ result = 1 * idx
1531
+ tm.assert_equal(result, idx)
1532
+
1533
+ def test_td64arr_mul_tdlike_scalar_raises(self, two_hours, box_with_array):
1534
+ rng = timedelta_range("1 days", "10 days", name="foo")
1535
+ rng = tm.box_expected(rng, box_with_array)
1536
+ msg = "|".join(
1537
+ [
1538
+ "argument must be an integer",
1539
+ "cannot use operands with types dtype",
1540
+ "Cannot multiply with",
1541
+ r"unsupported operand type\(s\) for \*",
1542
+ ]
1543
+ )
1544
+ with pytest.raises(TypeError, match=msg):
1545
+ rng * two_hours
1546
+
1547
+ def test_tdi_mul_int_array_zerodim(self, box_with_array):
1548
+ rng5 = np.arange(5, dtype="int64")
1549
+ idx = TimedeltaIndex(rng5)
1550
+ expected = TimedeltaIndex(rng5 * 5)
1551
+
1552
+ idx = tm.box_expected(idx, box_with_array)
1553
+ expected = tm.box_expected(expected, box_with_array)
1554
+
1555
+ result = idx * np.array(5, dtype="int64")
1556
+ tm.assert_equal(result, expected)
1557
+
1558
+ def test_tdi_mul_int_array(self, box_with_array):
1559
+ rng5 = np.arange(5, dtype="int64")
1560
+ idx = TimedeltaIndex(rng5)
1561
+ expected = TimedeltaIndex(rng5**2)
1562
+
1563
+ idx = tm.box_expected(idx, box_with_array)
1564
+ expected = tm.box_expected(expected, box_with_array)
1565
+
1566
+ result = idx * rng5
1567
+ tm.assert_equal(result, expected)
1568
+
1569
+ def test_tdi_mul_int_series(self, box_with_array):
1570
+ box = box_with_array
1571
+ xbox = Series if box in [Index, tm.to_array, pd.array] else box
1572
+
1573
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
1574
+ expected = TimedeltaIndex(np.arange(5, dtype="int64") ** 2)
1575
+
1576
+ idx = tm.box_expected(idx, box)
1577
+ expected = tm.box_expected(expected, xbox)
1578
+
1579
+ result = idx * Series(np.arange(5, dtype="int64"))
1580
+ tm.assert_equal(result, expected)
1581
+
1582
+ def test_tdi_mul_float_series(self, box_with_array):
1583
+ box = box_with_array
1584
+ xbox = Series if box in [Index, tm.to_array, pd.array] else box
1585
+
1586
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
1587
+ idx = tm.box_expected(idx, box)
1588
+
1589
+ rng5f = np.arange(5, dtype="float64")
1590
+ expected = TimedeltaIndex(rng5f * (rng5f + 1.0))
1591
+ expected = tm.box_expected(expected, xbox)
1592
+
1593
+ result = idx * Series(rng5f + 1.0)
1594
+ tm.assert_equal(result, expected)
1595
+
1596
+ # TODO: Put Series/DataFrame in others?
1597
+ @pytest.mark.parametrize(
1598
+ "other",
1599
+ [
1600
+ np.arange(1, 11),
1601
+ Index(np.arange(1, 11), np.int64),
1602
+ Index(range(1, 11), np.uint64),
1603
+ Index(range(1, 11), np.float64),
1604
+ pd.RangeIndex(1, 11),
1605
+ ],
1606
+ ids=lambda x: type(x).__name__,
1607
+ )
1608
+ def test_tdi_rmul_arraylike(self, other, box_with_array):
1609
+ box = box_with_array
1610
+
1611
+ tdi = TimedeltaIndex(["1 Day"] * 10)
1612
+ expected = timedelta_range("1 days", "10 days")._with_freq(None)
1613
+
1614
+ tdi = tm.box_expected(tdi, box)
1615
+ xbox = get_upcast_box(tdi, other)
1616
+
1617
+ expected = tm.box_expected(expected, xbox)
1618
+
1619
+ result = other * tdi
1620
+ tm.assert_equal(result, expected)
1621
+ commute = tdi * other
1622
+ tm.assert_equal(commute, expected)
1623
+
1624
+ def test_td64arr_mul_bool_scalar_raises(self, box_with_array):
1625
+ # GH#58054
1626
+ ser = Series(np.arange(5) * timedelta(hours=1), dtype="m8[ns]")
1627
+ obj = tm.box_expected(ser, box_with_array)
1628
+
1629
+ msg = r"Cannot multiply 'timedelta64\[ns\]' by bool"
1630
+ with pytest.raises(TypeError, match=msg):
1631
+ True * obj
1632
+ with pytest.raises(TypeError, match=msg):
1633
+ obj * True
1634
+ with pytest.raises(TypeError, match=msg):
1635
+ np.True_ * obj
1636
+ with pytest.raises(TypeError, match=msg):
1637
+ obj * np.True_
1638
+
1639
+ @pytest.mark.parametrize(
1640
+ "dtype",
1641
+ [
1642
+ bool,
1643
+ "boolean",
1644
+ pytest.param("bool[pyarrow]", marks=td.skip_if_no("pyarrow")),
1645
+ ],
1646
+ )
1647
+ def test_td64arr_mul_bool_raises(self, dtype, box_with_array):
1648
+ # GH#58054
1649
+ ser = Series(np.arange(5) * timedelta(hours=1), dtype="m8[ns]")
1650
+ obj = tm.box_expected(ser, box_with_array)
1651
+
1652
+ other = Series(np.arange(5) < 0.5, dtype=dtype)
1653
+ other = tm.box_expected(other, box_with_array)
1654
+
1655
+ msg = r"Cannot multiply 'timedelta64\[ns\]' by bool"
1656
+ with pytest.raises(TypeError, match=msg):
1657
+ obj * other
1658
+
1659
+ msg2 = msg.replace("rmul", "mul")
1660
+ if dtype == "bool[pyarrow]":
1661
+ # We go through ArrowEA.__mul__ which gives a different message
1662
+ msg2 = (
1663
+ r"operation 'mul' not supported for dtype 'bool\[pyarrow\]' "
1664
+ r"with dtype 'timedelta64\[ns\]'"
1665
+ )
1666
+ with pytest.raises(TypeError, match=msg2):
1667
+ other * obj
1668
+
1669
+ @pytest.mark.parametrize(
1670
+ "dtype",
1671
+ [
1672
+ "Int64",
1673
+ "Float64",
1674
+ pytest.param("int64[pyarrow]", marks=td.skip_if_no("pyarrow")),
1675
+ ],
1676
+ )
1677
+ def test_td64arr_mul_masked(self, dtype, box_with_array):
1678
+ ser = Series(np.arange(5) * timedelta(hours=1), dtype="m8[ns]")
1679
+ obj = tm.box_expected(ser, box_with_array)
1680
+
1681
+ other = Series(np.arange(5), dtype=dtype)
1682
+ other = tm.box_expected(other, box_with_array)
1683
+
1684
+ expected = Series([Timedelta(hours=n**2) for n in range(5)], dtype="m8[ns]")
1685
+ expected = tm.box_expected(expected, box_with_array)
1686
+ if dtype == "int64[pyarrow]":
1687
+ expected = expected.astype("duration[ns][pyarrow]")
1688
+
1689
+ result = obj * other
1690
+ tm.assert_equal(result, expected)
1691
+ result = other * obj
1692
+ tm.assert_equal(result, expected)
1693
+
1694
+ # ------------------------------------------------------------------
1695
+ # __div__, __rdiv__
1696
+
1697
+ def test_td64arr_div_nat_invalid(self, box_with_array):
1698
+ # don't allow division by NaT (maybe could in the future)
1699
+ rng = timedelta_range("1 days", "10 days", name="foo")
1700
+ rng = tm.box_expected(rng, box_with_array)
1701
+
1702
+ with pytest.raises(TypeError, match="unsupported operand type"):
1703
+ rng / NaT
1704
+ with pytest.raises(TypeError, match="Cannot divide NaTType by"):
1705
+ NaT / rng
1706
+
1707
+ dt64nat = np.datetime64("NaT", "ns")
1708
+ msg = "|".join(
1709
+ [
1710
+ # 'divide' on npdev as of 2021-12-18
1711
+ "ufunc '(true_divide|divide)' cannot use operands",
1712
+ "cannot perform __r?truediv__",
1713
+ "Cannot divide datetime64 by TimedeltaArray",
1714
+ ]
1715
+ )
1716
+ with pytest.raises(TypeError, match=msg):
1717
+ rng / dt64nat
1718
+ with pytest.raises(TypeError, match=msg):
1719
+ dt64nat / rng
1720
+
1721
+ def test_td64arr_div_td64nat(self, box_with_array):
1722
+ # GH#23829
1723
+ box = box_with_array
1724
+ xbox = np.ndarray if box is pd.array else box
1725
+
1726
+ rng = timedelta_range("1 days", "10 days")
1727
+ rng = tm.box_expected(rng, box)
1728
+
1729
+ other = np.timedelta64("NaT", "ns")
1730
+
1731
+ expected = np.array([np.nan] * 10)
1732
+ expected = tm.box_expected(expected, xbox)
1733
+
1734
+ result = rng / other
1735
+ tm.assert_equal(result, expected)
1736
+
1737
+ result = other / rng
1738
+ tm.assert_equal(result, expected)
1739
+
1740
+ def test_td64arr_div_int(self, box_with_array):
1741
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
1742
+ idx = tm.box_expected(idx, box_with_array)
1743
+
1744
+ result = idx / 1
1745
+ tm.assert_equal(result, idx)
1746
+
1747
+ with pytest.raises(TypeError, match="Cannot divide"):
1748
+ # GH#23829
1749
+ 1 / idx
1750
+
1751
+ def test_td64arr_div_tdlike_scalar(self, two_hours, box_with_array):
1752
+ # GH#20088, GH#22163 ensure DataFrame returns correct dtype
1753
+ box = box_with_array
1754
+ xbox = np.ndarray if box is pd.array else box
1755
+
1756
+ rng = timedelta_range("1 days", "10 days", name="foo")
1757
+ expected = Index((np.arange(10) + 1) * 12, dtype=np.float64, name="foo")
1758
+
1759
+ rng = tm.box_expected(rng, box)
1760
+ expected = tm.box_expected(expected, xbox)
1761
+
1762
+ result = rng / two_hours
1763
+ tm.assert_equal(result, expected)
1764
+
1765
+ result = two_hours / rng
1766
+ expected = 1 / expected
1767
+ tm.assert_equal(result, expected)
1768
+
1769
+ @pytest.mark.parametrize("m", [1, 3, 10])
1770
+ @pytest.mark.parametrize("unit", ["D", "h", "m", "s", "ms", "us", "ns"])
1771
+ def test_td64arr_div_td64_scalar(self, m, unit, box_with_array):
1772
+ box = box_with_array
1773
+ xbox = np.ndarray if box is pd.array else box
1774
+
1775
+ ser = Series([Timedelta(days=59)] * 3)
1776
+ ser[2] = np.nan
1777
+ flat = ser
1778
+ ser = tm.box_expected(ser, box)
1779
+
1780
+ # op
1781
+ expected = Series([x / np.timedelta64(m, unit) for x in flat])
1782
+ expected = tm.box_expected(expected, xbox)
1783
+ result = ser / np.timedelta64(m, unit)
1784
+ tm.assert_equal(result, expected)
1785
+
1786
+ # reverse op
1787
+ expected = Series([Timedelta(np.timedelta64(m, unit)) / x for x in flat])
1788
+ expected = tm.box_expected(expected, xbox)
1789
+ result = np.timedelta64(m, unit) / ser
1790
+ tm.assert_equal(result, expected)
1791
+
1792
+ def test_td64arr_div_tdlike_scalar_with_nat(self, two_hours, box_with_array):
1793
+ box = box_with_array
1794
+ xbox = np.ndarray if box is pd.array else box
1795
+
1796
+ rng = TimedeltaIndex(["1 days", NaT, "2 days"], name="foo")
1797
+ expected = Index([12, np.nan, 24], dtype=np.float64, name="foo")
1798
+
1799
+ rng = tm.box_expected(rng, box)
1800
+ expected = tm.box_expected(expected, xbox)
1801
+
1802
+ result = rng / two_hours
1803
+ tm.assert_equal(result, expected)
1804
+
1805
+ result = two_hours / rng
1806
+ expected = 1 / expected
1807
+ tm.assert_equal(result, expected)
1808
+
1809
+ def test_td64arr_div_td64_ndarray(self, box_with_array):
1810
+ # GH#22631
1811
+ box = box_with_array
1812
+ xbox = np.ndarray if box is pd.array else box
1813
+
1814
+ rng = TimedeltaIndex(["1 days", NaT, "2 days"])
1815
+ expected = Index([12, np.nan, 24], dtype=np.float64)
1816
+
1817
+ rng = tm.box_expected(rng, box)
1818
+ expected = tm.box_expected(expected, xbox)
1819
+
1820
+ other = np.array([2, 4, 2], dtype="m8[h]")
1821
+ result = rng / other
1822
+ tm.assert_equal(result, expected)
1823
+
1824
+ result = rng / tm.box_expected(other, box)
1825
+ tm.assert_equal(result, expected)
1826
+
1827
+ result = rng / other.astype(object)
1828
+ tm.assert_equal(result, expected.astype(object))
1829
+
1830
+ result = rng / list(other)
1831
+ tm.assert_equal(result, expected)
1832
+
1833
+ # reversed op
1834
+ expected = 1 / expected
1835
+ result = other / rng
1836
+ tm.assert_equal(result, expected)
1837
+
1838
+ result = tm.box_expected(other, box) / rng
1839
+ tm.assert_equal(result, expected)
1840
+
1841
+ result = other.astype(object) / rng
1842
+ tm.assert_equal(result, expected)
1843
+
1844
+ result = list(other) / rng
1845
+ tm.assert_equal(result, expected)
1846
+
1847
+ def test_tdarr_div_length_mismatch(self, box_with_array):
1848
+ rng = TimedeltaIndex(["1 days", NaT, "2 days"])
1849
+ mismatched = [1, 2, 3, 4]
1850
+
1851
+ rng = tm.box_expected(rng, box_with_array)
1852
+ msg = "Cannot divide vectors|Unable to coerce to Series"
1853
+ for obj in [mismatched, mismatched[:2]]:
1854
+ # one shorter, one longer
1855
+ for other in [obj, np.array(obj), Index(obj)]:
1856
+ with pytest.raises(ValueError, match=msg):
1857
+ rng / other
1858
+ with pytest.raises(ValueError, match=msg):
1859
+ other / rng
1860
+
1861
+ def test_td64_div_object_mixed_result(self, box_with_array):
1862
+ # Case where we having a NaT in the result inseat of timedelta64("NaT")
1863
+ # is misleading
1864
+ orig = timedelta_range("1 Day", periods=3).insert(1, NaT)
1865
+ tdi = tm.box_expected(orig, box_with_array, transpose=False)
1866
+
1867
+ other = np.array([orig[0], 1.5, 2.0, orig[2]], dtype=object)
1868
+ other = tm.box_expected(other, box_with_array, transpose=False)
1869
+
1870
+ res = tdi / other
1871
+
1872
+ expected = Index([1.0, np.timedelta64("NaT", "us"), orig[0], 1.5], dtype=object)
1873
+ expected = tm.box_expected(expected, box_with_array, transpose=False)
1874
+ if isinstance(expected, NumpyExtensionArray):
1875
+ expected = expected.to_numpy()
1876
+ tm.assert_equal(res, expected)
1877
+ if box_with_array is DataFrame:
1878
+ # We have an np.timedelta64(NaT), not pd.NaT
1879
+ assert isinstance(res.iloc[1, 0], np.timedelta64)
1880
+
1881
+ res = tdi // other
1882
+
1883
+ expected = Index([1, np.timedelta64("NaT", "us"), orig[0], 1], dtype=object)
1884
+ expected = tm.box_expected(expected, box_with_array, transpose=False)
1885
+ if isinstance(expected, NumpyExtensionArray):
1886
+ expected = expected.to_numpy()
1887
+ tm.assert_equal(res, expected)
1888
+ if box_with_array is DataFrame:
1889
+ # We have an np.timedelta64(NaT), not pd.NaT
1890
+ assert isinstance(res.iloc[1, 0], np.timedelta64)
1891
+
1892
+ # ------------------------------------------------------------------
1893
+ # __floordiv__, __rfloordiv__
1894
+
1895
+ @pytest.mark.skipif(WASM, reason="no fp exception support in wasm")
1896
+ def test_td64arr_floordiv_td64arr_with_nat(self, box_with_array):
1897
+ # GH#35529
1898
+ box = box_with_array
1899
+ xbox = np.ndarray if box is pd.array else box
1900
+
1901
+ left = Series([1000, 222330, 30], dtype="timedelta64[ns]")
1902
+ right = Series([1000, 222330, None], dtype="timedelta64[ns]")
1903
+
1904
+ left = tm.box_expected(left, box)
1905
+ right = tm.box_expected(right, box)
1906
+
1907
+ expected = np.array([1.0, 1.0, np.nan], dtype=np.float64)
1908
+ expected = tm.box_expected(expected, xbox)
1909
+
1910
+ with tm.maybe_produces_warning(
1911
+ RuntimeWarning, box is pd.array, check_stacklevel=False
1912
+ ):
1913
+ result = left // right
1914
+
1915
+ tm.assert_equal(result, expected)
1916
+
1917
+ # case that goes through __rfloordiv__ with arraylike
1918
+ with tm.maybe_produces_warning(
1919
+ RuntimeWarning, box is pd.array, check_stacklevel=False
1920
+ ):
1921
+ result = np.asarray(left) // right
1922
+ tm.assert_equal(result, expected)
1923
+
1924
+ @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning")
1925
+ def test_td64arr_floordiv_tdscalar(self, box_with_array, scalar_td):
1926
+ # GH#18831, GH#19125
1927
+ box = box_with_array
1928
+ xbox = np.ndarray if box is pd.array else box
1929
+ td = Timedelta("5m3s") # i.e. (scalar_td - 1sec) / 2
1930
+
1931
+ td1 = Series([td, td, NaT], dtype="m8[ns]")
1932
+ td1 = tm.box_expected(td1, box, transpose=False)
1933
+
1934
+ expected = Series([0, 0, np.nan])
1935
+ expected = tm.box_expected(expected, xbox, transpose=False)
1936
+
1937
+ result = td1 // scalar_td
1938
+ tm.assert_equal(result, expected)
1939
+
1940
+ # Reversed op
1941
+ expected = Series([2, 2, np.nan])
1942
+ expected = tm.box_expected(expected, xbox, transpose=False)
1943
+
1944
+ result = scalar_td // td1
1945
+ tm.assert_equal(result, expected)
1946
+
1947
+ # same thing buts let's be explicit about calling __rfloordiv__
1948
+ result = td1.__rfloordiv__(scalar_td)
1949
+ tm.assert_equal(result, expected)
1950
+
1951
+ def test_td64arr_floordiv_int(self, box_with_array):
1952
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
1953
+ idx = tm.box_expected(idx, box_with_array)
1954
+ result = idx // 1
1955
+ tm.assert_equal(result, idx)
1956
+
1957
+ pattern = "floor_divide cannot use operands|Cannot divide int by Timedelta*"
1958
+ with pytest.raises(TypeError, match=pattern):
1959
+ 1 // idx
1960
+
1961
+ # ------------------------------------------------------------------
1962
+ # mod, divmod
1963
+ # TODO: operations with timedelta-like arrays, numeric arrays,
1964
+ # reversed ops
1965
+
1966
+ def test_td64arr_mod_tdscalar(
1967
+ self, performance_warning, box_with_array, three_days
1968
+ ):
1969
+ tdi = timedelta_range("1 Day", "9 days")
1970
+ tdarr = tm.box_expected(tdi, box_with_array)
1971
+
1972
+ expected = TimedeltaIndex(["1 Day", "2 Days", "0 Days"] * 3)
1973
+ expected = tm.box_expected(expected, box_with_array)
1974
+
1975
+ if isinstance(three_days, offsets.Day):
1976
+ msg = "unsupported operand type"
1977
+ with pytest.raises(TypeError, match=msg):
1978
+ tdarr % three_days
1979
+ with pytest.raises(TypeError, match=msg):
1980
+ divmod(tdarr, three_days)
1981
+ with pytest.raises(TypeError, match=msg):
1982
+ tdarr // three_days
1983
+ return
1984
+
1985
+ result = tdarr % three_days
1986
+ tm.assert_equal(result, expected)
1987
+
1988
+ if box_with_array is DataFrame and isinstance(three_days, pd.DateOffset):
1989
+ # TODO: making expected be object here a result of DataFrame.__divmod__
1990
+ # being defined in a naive way that does not dispatch to the underlying
1991
+ # array's __divmod__
1992
+ expected = expected.astype(object)
1993
+ else:
1994
+ performance_warning = False
1995
+
1996
+ with tm.assert_produces_warning(performance_warning):
1997
+ result = divmod(tdarr, three_days)
1998
+
1999
+ tm.assert_equal(result[1], expected)
2000
+ tm.assert_equal(result[0], tdarr // three_days)
2001
+
2002
+ def test_td64arr_mod_int(self, box_with_array):
2003
+ tdi = timedelta_range("1 ns", "10 ns", periods=10)
2004
+ tdarr = tm.box_expected(tdi, box_with_array)
2005
+
2006
+ expected = TimedeltaIndex(["1 ns", "0 ns"] * 5)
2007
+ expected = tm.box_expected(expected, box_with_array)
2008
+
2009
+ result = tdarr % 2
2010
+ tm.assert_equal(result, expected)
2011
+
2012
+ msg = "Cannot divide int by"
2013
+ with pytest.raises(TypeError, match=msg):
2014
+ 2 % tdarr
2015
+
2016
+ result = divmod(tdarr, 2)
2017
+ tm.assert_equal(result[1], expected)
2018
+ tm.assert_equal(result[0], tdarr // 2)
2019
+
2020
+ def test_td64arr_rmod_tdscalar(self, box_with_array, three_days):
2021
+ tdi = timedelta_range("1 Day", "9 days")
2022
+ tdarr = tm.box_expected(tdi, box_with_array)
2023
+
2024
+ expected = ["0 Days", "1 Day", "0 Days"] + ["3 Days"] * 6
2025
+ expected = TimedeltaIndex(expected)
2026
+ expected = tm.box_expected(expected, box_with_array)
2027
+
2028
+ if isinstance(three_days, offsets.Day):
2029
+ msg = "Cannot divide Day by TimedeltaArray"
2030
+ with pytest.raises(TypeError, match=msg):
2031
+ three_days % tdarr
2032
+ return
2033
+
2034
+ result = three_days % tdarr
2035
+ tm.assert_equal(result, expected)
2036
+
2037
+ result = divmod(three_days, tdarr)
2038
+ tm.assert_equal(result[1], expected)
2039
+ tm.assert_equal(result[0], three_days // tdarr)
2040
+
2041
+ # ------------------------------------------------------------------
2042
+ # Operations with invalid others
2043
+
2044
+ def test_td64arr_mul_tdscalar_invalid(self, box_with_array, scalar_td):
2045
+ td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
2046
+ td1.iloc[2] = np.nan
2047
+
2048
+ td1 = tm.box_expected(td1, box_with_array)
2049
+
2050
+ # check that we are getting a TypeError
2051
+ # with 'operate' (from core/ops.py) for the ops that are not
2052
+ # defined
2053
+ pattern = "operate|unsupported|cannot|not supported"
2054
+ with pytest.raises(TypeError, match=pattern):
2055
+ td1 * scalar_td
2056
+ with pytest.raises(TypeError, match=pattern):
2057
+ scalar_td * td1
2058
+
2059
+ def test_td64arr_mul_too_short_raises(self, box_with_array):
2060
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
2061
+ idx = tm.box_expected(idx, box_with_array)
2062
+ msg = "|".join(
2063
+ [
2064
+ "cannot use operands with types dtype",
2065
+ "Cannot multiply with unequal lengths",
2066
+ "Unable to coerce to Series",
2067
+ ]
2068
+ )
2069
+ with pytest.raises(TypeError, match=msg):
2070
+ # length check before dtype check
2071
+ idx * idx[:3]
2072
+ with pytest.raises(ValueError, match=msg):
2073
+ idx * np.array([1, 2])
2074
+
2075
+ def test_td64arr_mul_td64arr_raises(self, box_with_array):
2076
+ idx = TimedeltaIndex(np.arange(5, dtype="int64"))
2077
+ idx = tm.box_expected(idx, box_with_array)
2078
+ msg = "cannot use operands with types dtype"
2079
+ with pytest.raises(TypeError, match=msg):
2080
+ idx * idx
2081
+
2082
+ # ------------------------------------------------------------------
2083
+ # Operations with numeric others
2084
+
2085
+ def test_td64arr_mul_numeric_scalar(self, box_with_array, one):
2086
+ # GH#4521
2087
+ # divide/multiply by integers
2088
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
2089
+ expected = Series(["-59 Days", "-59 Days", "NaT"], dtype="timedelta64[ns]")
2090
+
2091
+ tdser = tm.box_expected(tdser, box_with_array)
2092
+ expected = tm.box_expected(expected, box_with_array)
2093
+
2094
+ result = tdser * (-one)
2095
+ tm.assert_equal(result, expected)
2096
+ result = (-one) * tdser
2097
+ tm.assert_equal(result, expected)
2098
+
2099
+ expected = Series(["118 Days", "118 Days", "NaT"], dtype="timedelta64[ns]")
2100
+ expected = tm.box_expected(expected, box_with_array)
2101
+
2102
+ result = tdser * (2 * one)
2103
+ tm.assert_equal(result, expected)
2104
+ result = (2 * one) * tdser
2105
+ tm.assert_equal(result, expected)
2106
+
2107
+ @pytest.mark.parametrize("two", [2, 2.0, np.array(2), np.array(2.0)])
2108
+ def test_td64arr_div_numeric_scalar(self, box_with_array, two):
2109
+ # GH#4521
2110
+ # divide/multiply by integers
2111
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
2112
+ expected = Series(["29.5D", "29.5D", "NaT"], dtype="timedelta64[ns]")
2113
+
2114
+ tdser = tm.box_expected(tdser, box_with_array)
2115
+ expected = tm.box_expected(expected, box_with_array)
2116
+
2117
+ result = tdser / two
2118
+ tm.assert_equal(result, expected)
2119
+
2120
+ with pytest.raises(TypeError, match="Cannot divide"):
2121
+ two / tdser
2122
+
2123
+ @pytest.mark.parametrize("two", [2, 2.0, np.array(2), np.array(2.0)])
2124
+ def test_td64arr_floordiv_numeric_scalar(self, box_with_array, two):
2125
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
2126
+ expected = Series(["29.5D", "29.5D", "NaT"], dtype="timedelta64[ns]")
2127
+
2128
+ tdser = tm.box_expected(tdser, box_with_array)
2129
+ expected = tm.box_expected(expected, box_with_array)
2130
+
2131
+ result = tdser // two
2132
+ tm.assert_equal(result, expected)
2133
+
2134
+ with pytest.raises(TypeError, match="Cannot divide"):
2135
+ two // tdser
2136
+
2137
+ @pytest.mark.parametrize(
2138
+ "klass",
2139
+ [np.array, Index, Series],
2140
+ ids=lambda x: x.__name__,
2141
+ )
2142
+ def test_td64arr_rmul_numeric_array(
2143
+ self,
2144
+ box_with_array,
2145
+ klass,
2146
+ any_real_numpy_dtype,
2147
+ ):
2148
+ # GH#4521
2149
+ # divide/multiply by integers
2150
+
2151
+ vector = klass([20, 30, 40])
2152
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
2153
+ vector = vector.astype(any_real_numpy_dtype)
2154
+
2155
+ expected = Series(["1180 Days", "1770 Days", "NaT"], dtype="timedelta64[ns]")
2156
+
2157
+ tdser = tm.box_expected(tdser, box_with_array)
2158
+ xbox = get_upcast_box(tdser, vector)
2159
+
2160
+ expected = tm.box_expected(expected, xbox)
2161
+
2162
+ result = tdser * vector
2163
+ tm.assert_equal(result, expected)
2164
+
2165
+ result = vector * tdser
2166
+ tm.assert_equal(result, expected)
2167
+
2168
+ @pytest.mark.parametrize(
2169
+ "klass",
2170
+ [np.array, Index, Series],
2171
+ ids=lambda x: x.__name__,
2172
+ )
2173
+ def test_td64arr_div_numeric_array(
2174
+ self, box_with_array, klass, any_real_numpy_dtype
2175
+ ):
2176
+ # GH#4521
2177
+ # divide/multiply by integers
2178
+
2179
+ vector = klass([20, 30, 40])
2180
+ tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]")
2181
+ vector = vector.astype(any_real_numpy_dtype)
2182
+
2183
+ expected = Series(["2.95D", "1D 23h 12m", "NaT"], dtype="timedelta64[ns]")
2184
+
2185
+ tdser = tm.box_expected(tdser, box_with_array)
2186
+ xbox = get_upcast_box(tdser, vector)
2187
+ expected = tm.box_expected(expected, xbox)
2188
+
2189
+ result = tdser / vector
2190
+ tm.assert_equal(result, expected)
2191
+
2192
+ pattern = "|".join(
2193
+ [
2194
+ "true_divide'? cannot use operands",
2195
+ "cannot perform __div__",
2196
+ "cannot perform __truediv__",
2197
+ "unsupported operand",
2198
+ "Cannot divide",
2199
+ "ufunc 'divide' cannot use operands with types",
2200
+ ]
2201
+ )
2202
+ with pytest.raises(TypeError, match=pattern):
2203
+ vector / tdser
2204
+
2205
+ result = tdser / vector.astype(object)
2206
+ if box_with_array is DataFrame:
2207
+ expected = [tdser.iloc[0, n] / vector[n] for n in range(len(vector))]
2208
+ expected = tm.box_expected(expected, xbox).astype(object)
2209
+ expected[2] = expected[2].fillna(np.timedelta64("NaT", "ns"))
2210
+ else:
2211
+ expected = [tdser[n] / vector[n] for n in range(len(tdser))]
2212
+ expected = [
2213
+ x if x is not NaT else np.timedelta64("NaT", "ns") for x in expected
2214
+ ]
2215
+ if xbox is tm.to_array:
2216
+ expected = tm.to_array(expected).astype(object)
2217
+ else:
2218
+ expected = xbox(expected, dtype=object)
2219
+
2220
+ tm.assert_equal(result, expected)
2221
+
2222
+ with pytest.raises(TypeError, match=pattern):
2223
+ vector.astype(object) / tdser
2224
+
2225
+ def test_td64arr_mul_int_series(self, box_with_array, names):
2226
+ # GH#19042 test for correct name attachment
2227
+ box = box_with_array
2228
+ exname = get_expected_name(box, names)
2229
+
2230
+ tdi = TimedeltaIndex(
2231
+ ["0days", "1day", "2days", "3days", "4days"], name=names[0]
2232
+ )
2233
+ # TODO: Should we be parametrizing over types for `ser` too?
2234
+ ser = Series([0, 1, 2, 3, 4], dtype=np.int64, name=names[1])
2235
+
2236
+ expected = Series(
2237
+ ["0days", "1day", "4days", "9days", "16days"],
2238
+ dtype="timedelta64[us]",
2239
+ name=exname,
2240
+ )
2241
+
2242
+ tdi = tm.box_expected(tdi, box)
2243
+ xbox = get_upcast_box(tdi, ser)
2244
+
2245
+ expected = tm.box_expected(expected, xbox)
2246
+
2247
+ result = ser * tdi
2248
+ tm.assert_equal(result, expected)
2249
+
2250
+ result = tdi * ser
2251
+ tm.assert_equal(result, expected)
2252
+
2253
+ # TODO: Should we be parametrizing over types for `ser` too?
2254
+ def test_float_series_rdiv_td64arr(self, box_with_array, names):
2255
+ # GH#19042 test for correct name attachment
2256
+ box = box_with_array
2257
+ tdi = TimedeltaIndex(
2258
+ ["0days", "1day", "2days", "3days", "4days"], name=names[0]
2259
+ )
2260
+ ser = Series([1.5, 3, 4.5, 6, 7.5], dtype=np.float64, name=names[1])
2261
+
2262
+ xname = names[2] if box not in [tm.to_array, pd.array] else names[1]
2263
+ expected = Series(
2264
+ [tdi[n] / ser[n] for n in range(len(ser))],
2265
+ dtype="timedelta64[us]",
2266
+ name=xname,
2267
+ )
2268
+
2269
+ tdi = tm.box_expected(tdi, box)
2270
+ xbox = get_upcast_box(tdi, ser)
2271
+ expected = tm.box_expected(expected, xbox)
2272
+
2273
+ result = ser.__rtruediv__(tdi)
2274
+ if box is DataFrame:
2275
+ assert result is NotImplemented
2276
+ else:
2277
+ tm.assert_equal(result, expected)
2278
+
2279
+ def test_td64arr_all_nat_div_object_dtype_numeric(self, box_with_array):
2280
+ # GH#39750 make sure we infer the result as td64
2281
+ tdi = TimedeltaIndex([NaT, NaT], dtype="m8[ns]")
2282
+
2283
+ left = tm.box_expected(tdi, box_with_array)
2284
+ right = np.array([2, 2.0], dtype=object)
2285
+
2286
+ tdnat = np.timedelta64("NaT", "ns")
2287
+ expected = Index([tdnat] * 2, dtype=object)
2288
+ if box_with_array is not Index:
2289
+ expected = tm.box_expected(expected, box_with_array).astype(object)
2290
+ if box_with_array in [Series, DataFrame]:
2291
+ expected = expected.fillna(tdnat) # GH#18463
2292
+
2293
+ result = left / right
2294
+ tm.assert_equal(result, expected)
2295
+
2296
+ result = left // right
2297
+ tm.assert_equal(result, expected)
2298
+
2299
+
2300
+ class TestTimedelta64ArrayLikeArithmetic:
2301
+ # Arithmetic tests for timedelta64[ns] vectors fully parametrized over
2302
+ # DataFrame/Series/TimedeltaIndex/TimedeltaArray. Ideally all arithmetic
2303
+ # tests will eventually end up here.
2304
+
2305
+ def test_td64arr_pow_invalid(self, scalar_td, box_with_array):
2306
+ td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
2307
+ td1.iloc[2] = np.nan
2308
+
2309
+ td1 = tm.box_expected(td1, box_with_array)
2310
+
2311
+ # check that we are getting a TypeError
2312
+ # with 'operate' (from core/ops.py) for the ops that are not
2313
+ # defined
2314
+ pattern = "operate|unsupported|cannot|not supported"
2315
+ with pytest.raises(TypeError, match=pattern):
2316
+ scalar_td**td1
2317
+
2318
+ with pytest.raises(TypeError, match=pattern):
2319
+ td1**scalar_td
2320
+
2321
+
2322
+ def test_add_timestamp_to_timedelta():
2323
+ # GH: 35897
2324
+ timestamp = Timestamp("2021-01-01")
2325
+ result = timestamp + timedelta_range("0s", "1s", periods=31)
2326
+ expected = DatetimeIndex(
2327
+ [
2328
+ timestamp
2329
+ + (
2330
+ pd.to_timedelta("0.033333s") * i
2331
+ + pd.to_timedelta("0.000001s") * divmod(i, 3)[0]
2332
+ )
2333
+ for i in range(31)
2334
+ ]
2335
+ )
2336
+ tm.assert_index_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/__init__.py ADDED
File without changes
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/__init__.py ADDED
File without changes
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_arithmetic.py ADDED
@@ -0,0 +1,134 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import operator
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ import pandas as pd
7
+ import pandas._testing as tm
8
+
9
+
10
+ @pytest.fixture
11
+ def data():
12
+ """Fixture returning boolean array with valid and missing values."""
13
+ return pd.array(
14
+ [True, False] * 2 + [np.nan] + [True, False] + [np.nan] + [True, False],
15
+ dtype="boolean",
16
+ )
17
+
18
+
19
+ @pytest.fixture
20
+ def left_array():
21
+ """Fixture returning boolean array with valid and missing values."""
22
+ return pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
23
+
24
+
25
+ @pytest.fixture
26
+ def right_array():
27
+ """Fixture returning boolean array with valid and missing values."""
28
+ return pd.array([True, False, None] * 3, dtype="boolean")
29
+
30
+
31
+ # Basic test for the arithmetic array ops
32
+ # -----------------------------------------------------------------------------
33
+
34
+
35
+ @pytest.mark.parametrize(
36
+ "opname, exp",
37
+ [
38
+ ("add", [True, True, None, True, False, None, None, None, None]),
39
+ ("mul", [True, False, None, False, False, None, None, None, None]),
40
+ ],
41
+ ids=["add", "mul"],
42
+ )
43
+ def test_add_mul(left_array, right_array, opname, exp):
44
+ op = getattr(operator, opname)
45
+ result = op(left_array, right_array)
46
+ expected = pd.array(exp, dtype="boolean")
47
+ tm.assert_extension_array_equal(result, expected)
48
+
49
+
50
+ def test_sub(left_array, right_array):
51
+ msg = (
52
+ r"numpy boolean subtract, the `-` operator, is (?:deprecated|not supported), "
53
+ r"use the bitwise_xor, the `\^` operator, or the logical_xor function instead\."
54
+ )
55
+ with pytest.raises(TypeError, match=msg):
56
+ left_array - right_array
57
+
58
+
59
+ def test_div(left_array, right_array):
60
+ msg = "operator '.*' not implemented for bool dtypes"
61
+ with pytest.raises(NotImplementedError, match=msg):
62
+ # check that we are matching the non-masked Series behavior
63
+ pd.Series(left_array._data) / pd.Series(right_array._data)
64
+
65
+ with pytest.raises(NotImplementedError, match=msg):
66
+ left_array / right_array
67
+
68
+
69
+ @pytest.mark.parametrize(
70
+ "opname",
71
+ [
72
+ "floordiv",
73
+ "mod",
74
+ "pow",
75
+ ],
76
+ )
77
+ def test_op_int8(left_array, right_array, opname):
78
+ op = getattr(operator, opname)
79
+ if opname != "mod":
80
+ msg = "operator '.*' not implemented for bool dtypes"
81
+ with pytest.raises(NotImplementedError, match=msg):
82
+ result = op(left_array, right_array)
83
+ return
84
+ result = op(left_array, right_array)
85
+ expected = op(left_array.astype("Int8"), right_array.astype("Int8"))
86
+ tm.assert_extension_array_equal(result, expected)
87
+
88
+
89
+ # Test generic characteristics / errors
90
+ # -----------------------------------------------------------------------------
91
+
92
+
93
+ def test_error_invalid_values(data, all_arithmetic_operators):
94
+ # invalid ops
95
+ op = all_arithmetic_operators
96
+ s = pd.Series(data)
97
+ ops = getattr(s, op)
98
+
99
+ # invalid scalars
100
+ msg = (
101
+ "did not contain a loop with signature matching types|"
102
+ "BooleanArray cannot perform the operation|"
103
+ "not supported for the input types, and the inputs could not be safely coerced "
104
+ "to any supported types according to the casting rule ''safe''|"
105
+ "not supported for dtype"
106
+ )
107
+ with pytest.raises(TypeError, match=msg):
108
+ ops("foo")
109
+ msg = "|".join(
110
+ [
111
+ r"unsupported operand type\(s\) for",
112
+ "Concatenation operation is not implemented for NumPy arrays",
113
+ "has no kernel",
114
+ "not supported for dtype",
115
+ ]
116
+ )
117
+ with pytest.raises(TypeError, match=msg):
118
+ ops(pd.Timestamp("20180101"))
119
+
120
+ # invalid array-likes
121
+ if op not in ("__mul__", "__rmul__"):
122
+ # TODO(extension) numpy's mul with object array sees booleans as numbers
123
+ msg = "|".join(
124
+ [
125
+ r"unsupported operand type\(s\) for",
126
+ "can only concatenate str",
127
+ "not all arguments converted during string formatting",
128
+ "has no kernel",
129
+ "not implemented",
130
+ "not supported for dtype",
131
+ ]
132
+ )
133
+ with pytest.raises(TypeError, match=msg):
134
+ ops(pd.Series("foo", index=s.index))
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_astype.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+
7
+
8
+ def test_astype(using_infer_string):
9
+ # with missing values
10
+ arr = pd.array([True, False, None], dtype="boolean")
11
+
12
+ with pytest.raises(ValueError, match="cannot convert NA to integer"):
13
+ arr.astype("int64")
14
+
15
+ with pytest.raises(ValueError, match="cannot convert float NaN to"):
16
+ arr.astype("bool")
17
+
18
+ result = arr.astype("float64")
19
+ expected = np.array([1, 0, np.nan], dtype="float64")
20
+ tm.assert_numpy_array_equal(result, expected)
21
+
22
+ result = arr.astype("str")
23
+ if using_infer_string:
24
+ expected = pd.array(
25
+ ["True", "False", None], dtype=pd.StringDtype(na_value=np.nan)
26
+ )
27
+ tm.assert_extension_array_equal(result, expected)
28
+ else:
29
+ expected = np.array(["True", "False", "<NA>"], dtype=f"{tm.ENDIAN}U5")
30
+ tm.assert_numpy_array_equal(result, expected)
31
+
32
+ # no missing values
33
+ arr = pd.array([True, False, True], dtype="boolean")
34
+ result = arr.astype("int64")
35
+ expected = np.array([1, 0, 1], dtype="int64")
36
+ tm.assert_numpy_array_equal(result, expected)
37
+
38
+ result = arr.astype("bool")
39
+ expected = np.array([True, False, True], dtype="bool")
40
+ tm.assert_numpy_array_equal(result, expected)
41
+
42
+
43
+ def test_astype_to_boolean_array():
44
+ # astype to BooleanArray
45
+ arr = pd.array([True, False, None], dtype="boolean")
46
+
47
+ result = arr.astype("boolean")
48
+ tm.assert_extension_array_equal(result, arr)
49
+ result = arr.astype(pd.BooleanDtype())
50
+ tm.assert_extension_array_equal(result, arr)
51
+
52
+
53
+ def test_astype_to_integer_array():
54
+ # astype to IntegerArray
55
+ arr = pd.array([True, False, None], dtype="boolean")
56
+
57
+ result = arr.astype("Int64")
58
+ expected = pd.array([1, 0, None], dtype="Int64")
59
+ tm.assert_extension_array_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_comparison.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+ from pandas.arrays import BooleanArray
7
+ from pandas.tests.arrays.masked_shared import ComparisonOps
8
+
9
+
10
+ @pytest.fixture
11
+ def data():
12
+ """Fixture returning boolean array with valid and missing data"""
13
+ return pd.array(
14
+ [True, False] * 2 + [np.nan] + [True, False] + [np.nan] + [True, False],
15
+ dtype="boolean",
16
+ )
17
+
18
+
19
+ @pytest.fixture
20
+ def dtype():
21
+ """Fixture returning BooleanDtype"""
22
+ return pd.BooleanDtype()
23
+
24
+
25
+ class TestComparisonOps(ComparisonOps):
26
+ def test_compare_scalar(self, data, comparison_op):
27
+ self._compare_other(data, comparison_op, True)
28
+
29
+ def test_compare_array(self, data, comparison_op):
30
+ other = pd.array([True] * len(data), dtype="boolean")
31
+ self._compare_other(data, comparison_op, other)
32
+ other = np.array([True] * len(data))
33
+ self._compare_other(data, comparison_op, other)
34
+ other = pd.Series([True] * len(data))
35
+ self._compare_other(data, comparison_op, other)
36
+
37
+ @pytest.mark.parametrize("other", [True, False, pd.NA])
38
+ def test_scalar(self, other, comparison_op, dtype):
39
+ ComparisonOps.test_scalar(self, other, comparison_op, dtype)
40
+
41
+ def test_array(self, comparison_op):
42
+ op = comparison_op
43
+ a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
44
+ b = pd.array([True, False, None] * 3, dtype="boolean")
45
+
46
+ result = op(a, b)
47
+
48
+ values = op(a._data, b._data)
49
+ mask = a._mask | b._mask
50
+ expected = BooleanArray(values, mask)
51
+ tm.assert_extension_array_equal(result, expected)
52
+
53
+ # ensure we haven't mutated anything inplace
54
+ result[0] = None
55
+ tm.assert_extension_array_equal(
56
+ a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
57
+ )
58
+ tm.assert_extension_array_equal(
59
+ b, pd.array([True, False, None] * 3, dtype="boolean")
60
+ )
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_construction.py ADDED
@@ -0,0 +1,335 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+ from pandas.arrays import BooleanArray
7
+ from pandas.core.arrays.boolean import coerce_to_array
8
+
9
+
10
+ def test_boolean_array_constructor():
11
+ values = np.array([True, False, True, False], dtype="bool")
12
+ mask = np.array([False, False, False, True], dtype="bool")
13
+
14
+ result = BooleanArray(values, mask)
15
+ expected = pd.array([True, False, True, None], dtype="boolean")
16
+ tm.assert_extension_array_equal(result, expected)
17
+
18
+ with pytest.raises(TypeError, match="values should be boolean numpy array"):
19
+ BooleanArray(values.tolist(), mask)
20
+
21
+ with pytest.raises(TypeError, match="mask should be boolean numpy array"):
22
+ BooleanArray(values, mask.tolist())
23
+
24
+ with pytest.raises(TypeError, match="values should be boolean numpy array"):
25
+ BooleanArray(values.astype(int), mask)
26
+
27
+ with pytest.raises(TypeError, match="mask should be boolean numpy array"):
28
+ BooleanArray(values, None)
29
+
30
+ with pytest.raises(ValueError, match="values.shape must match mask.shape"):
31
+ BooleanArray(values.reshape(1, -1), mask)
32
+
33
+ with pytest.raises(ValueError, match="values.shape must match mask.shape"):
34
+ BooleanArray(values, mask.reshape(1, -1))
35
+
36
+
37
+ def test_boolean_array_constructor_copy():
38
+ values = np.array([True, False, True, False], dtype="bool")
39
+ mask = np.array([False, False, False, True], dtype="bool")
40
+
41
+ result = BooleanArray(values, mask)
42
+ assert result._data is values
43
+ assert result._mask is mask
44
+
45
+ result = BooleanArray(values, mask, copy=True)
46
+ assert result._data is not values
47
+ assert result._mask is not mask
48
+
49
+
50
+ def test_to_boolean_array():
51
+ expected = BooleanArray(
52
+ np.array([True, False, True]), np.array([False, False, False])
53
+ )
54
+
55
+ result = pd.array([True, False, True], dtype="boolean")
56
+ tm.assert_extension_array_equal(result, expected)
57
+ result = pd.array(np.array([True, False, True]), dtype="boolean")
58
+ tm.assert_extension_array_equal(result, expected)
59
+ result = pd.array(np.array([True, False, True], dtype=object), dtype="boolean")
60
+ tm.assert_extension_array_equal(result, expected)
61
+
62
+ # with missing values
63
+ expected = BooleanArray(
64
+ np.array([True, False, True]), np.array([False, False, True])
65
+ )
66
+
67
+ result = pd.array([True, False, None], dtype="boolean")
68
+ tm.assert_extension_array_equal(result, expected)
69
+ result = pd.array(np.array([True, False, None], dtype=object), dtype="boolean")
70
+ tm.assert_extension_array_equal(result, expected)
71
+
72
+
73
+ def test_to_boolean_array_all_none():
74
+ expected = BooleanArray(np.array([True, True, True]), np.array([True, True, True]))
75
+
76
+ result = pd.array([None, None, None], dtype="boolean")
77
+ tm.assert_extension_array_equal(result, expected)
78
+ result = pd.array(np.array([None, None, None], dtype=object), dtype="boolean")
79
+ tm.assert_extension_array_equal(result, expected)
80
+
81
+
82
+ @pytest.mark.parametrize(
83
+ "a, b",
84
+ [
85
+ ([True, False, None, np.nan, pd.NA], [True, False, None, None, None]),
86
+ ([True, np.nan], [True, None]),
87
+ ([True, pd.NA], [True, None]),
88
+ ([np.nan, np.nan], [None, None]),
89
+ (np.array([np.nan, np.nan], dtype=float), [None, None]),
90
+ ],
91
+ )
92
+ def test_to_boolean_array_missing_indicators(a, b):
93
+ result = pd.array(a, dtype="boolean")
94
+ expected = pd.array(b, dtype="boolean")
95
+ tm.assert_extension_array_equal(result, expected)
96
+
97
+
98
+ @pytest.mark.parametrize(
99
+ "values",
100
+ [
101
+ ["foo", "bar"],
102
+ ["1", "2"],
103
+ # "foo",
104
+ [1, 2],
105
+ [1.0, 2.0],
106
+ pd.date_range("20130101", periods=2),
107
+ np.array(["foo"]),
108
+ np.array([1, 2]),
109
+ np.array([1.0, 2.0]),
110
+ [np.nan, {"a": 1}],
111
+ ],
112
+ )
113
+ def test_to_boolean_array_error(values):
114
+ # error in converting existing arrays to BooleanArray
115
+ msg = "Need to pass bool-like value"
116
+ with pytest.raises(TypeError, match=msg):
117
+ pd.array(values, dtype="boolean")
118
+
119
+
120
+ def test_to_boolean_array_from_integer_array():
121
+ result = pd.array(np.array([1, 0, 1, 0]), dtype="boolean")
122
+ expected = pd.array([True, False, True, False], dtype="boolean")
123
+ tm.assert_extension_array_equal(result, expected)
124
+
125
+ # with missing values
126
+ result = pd.array(np.array([1, 0, 1, None]), dtype="boolean")
127
+ expected = pd.array([True, False, True, None], dtype="boolean")
128
+ tm.assert_extension_array_equal(result, expected)
129
+
130
+
131
+ def test_to_boolean_array_from_float_array():
132
+ result = pd.array(np.array([1.0, 0.0, 1.0, 0.0]), dtype="boolean")
133
+ expected = pd.array([True, False, True, False], dtype="boolean")
134
+ tm.assert_extension_array_equal(result, expected)
135
+
136
+ # with missing values
137
+ result = pd.array(np.array([1.0, 0.0, 1.0, np.nan]), dtype="boolean")
138
+ expected = pd.array([True, False, True, None], dtype="boolean")
139
+ tm.assert_extension_array_equal(result, expected)
140
+
141
+
142
+ def test_to_boolean_array_integer_like():
143
+ # integers of 0's and 1's
144
+ result = pd.array([1, 0, 1, 0], dtype="boolean")
145
+ expected = pd.array([True, False, True, False], dtype="boolean")
146
+ tm.assert_extension_array_equal(result, expected)
147
+
148
+ # with missing values
149
+ result = pd.array([1, 0, 1, None], dtype="boolean")
150
+ expected = pd.array([True, False, True, None], dtype="boolean")
151
+ tm.assert_extension_array_equal(result, expected)
152
+
153
+
154
+ def test_coerce_to_array():
155
+ # TODO this is currently not public API
156
+ values = np.array([True, False, True, False], dtype="bool")
157
+ mask = np.array([False, False, False, True], dtype="bool")
158
+ result = BooleanArray(*coerce_to_array(values, mask=mask))
159
+ expected = BooleanArray(values, mask)
160
+ tm.assert_extension_array_equal(result, expected)
161
+ assert result._data is values
162
+ assert result._mask is mask
163
+ result = BooleanArray(*coerce_to_array(values, mask=mask, copy=True))
164
+ expected = BooleanArray(values, mask)
165
+ tm.assert_extension_array_equal(result, expected)
166
+ assert result._data is not values
167
+ assert result._mask is not mask
168
+
169
+ # mixed missing from values and mask
170
+ values = [True, False, None, False]
171
+ mask = np.array([False, False, False, True], dtype="bool")
172
+ result = BooleanArray(*coerce_to_array(values, mask=mask))
173
+ expected = BooleanArray(
174
+ np.array([True, False, True, True]), np.array([False, False, True, True])
175
+ )
176
+ tm.assert_extension_array_equal(result, expected)
177
+ result = BooleanArray(*coerce_to_array(np.array(values, dtype=object), mask=mask))
178
+ tm.assert_extension_array_equal(result, expected)
179
+ result = BooleanArray(*coerce_to_array(values, mask=mask.tolist()))
180
+ tm.assert_extension_array_equal(result, expected)
181
+
182
+ # raise errors for wrong dimension
183
+ values = np.array([True, False, True, False], dtype="bool")
184
+ mask = np.array([False, False, False, True], dtype="bool")
185
+
186
+ # passing 2D values is OK as long as no mask
187
+ coerce_to_array(values.reshape(1, -1))
188
+
189
+ with pytest.raises(ValueError, match="values.shape and mask.shape must match"):
190
+ coerce_to_array(values.reshape(1, -1), mask=mask)
191
+
192
+ with pytest.raises(ValueError, match="values.shape and mask.shape must match"):
193
+ coerce_to_array(values, mask=mask.reshape(1, -1))
194
+
195
+
196
+ def test_coerce_to_array_from_boolean_array():
197
+ # passing BooleanArray to coerce_to_array
198
+ values = np.array([True, False, True, False], dtype="bool")
199
+ mask = np.array([False, False, False, True], dtype="bool")
200
+ arr = BooleanArray(values, mask)
201
+ result = BooleanArray(*coerce_to_array(arr))
202
+ tm.assert_extension_array_equal(result, arr)
203
+ # no copy
204
+ assert result._data is arr._data
205
+ assert result._mask is arr._mask
206
+
207
+ result = BooleanArray(*coerce_to_array(arr), copy=True)
208
+ tm.assert_extension_array_equal(result, arr)
209
+ assert result._data is not arr._data
210
+ assert result._mask is not arr._mask
211
+
212
+ with pytest.raises(ValueError, match="cannot pass mask for BooleanArray input"):
213
+ coerce_to_array(arr, mask=mask)
214
+
215
+
216
+ def test_coerce_to_numpy_array():
217
+ # with missing values -> object dtype
218
+ arr = pd.array([True, False, None], dtype="boolean")
219
+ result = np.array(arr)
220
+ expected = np.array([True, False, pd.NA], dtype="object")
221
+ tm.assert_numpy_array_equal(result, expected)
222
+
223
+ # also with no missing values -> object dtype
224
+ arr = pd.array([True, False, True], dtype="boolean")
225
+ result = np.array(arr)
226
+ expected = np.array([True, False, True], dtype="bool")
227
+ tm.assert_numpy_array_equal(result, expected)
228
+
229
+ # force bool dtype
230
+ result = np.array(arr, dtype="bool")
231
+ expected = np.array([True, False, True], dtype="bool")
232
+ tm.assert_numpy_array_equal(result, expected)
233
+ # with missing values will raise error
234
+ arr = pd.array([True, False, None], dtype="boolean")
235
+ msg = (
236
+ "cannot convert to 'bool'-dtype NumPy array with missing values. "
237
+ "Specify an appropriate 'na_value' for this dtype."
238
+ )
239
+ with pytest.raises(ValueError, match=msg):
240
+ np.array(arr, dtype="bool")
241
+
242
+
243
+ def test_to_boolean_array_from_strings():
244
+ result = BooleanArray._from_sequence_of_strings(
245
+ np.array(["True", "False", "1", "1.0", "0", "0.0", np.nan], dtype=object),
246
+ dtype=pd.BooleanDtype(),
247
+ )
248
+ expected = BooleanArray(
249
+ np.array([True, False, True, True, False, False, False]),
250
+ np.array([False, False, False, False, False, False, True]),
251
+ )
252
+
253
+ tm.assert_extension_array_equal(result, expected)
254
+
255
+
256
+ def test_to_boolean_array_from_strings_invalid_string():
257
+ with pytest.raises(ValueError, match="cannot be cast"):
258
+ BooleanArray._from_sequence_of_strings(["donkey"], dtype=pd.BooleanDtype())
259
+
260
+
261
+ @pytest.mark.parametrize("box", [True, False], ids=["series", "array"])
262
+ def test_to_numpy(box):
263
+ con = pd.Series if box else pd.array
264
+ # default (with or without missing values) -> object dtype
265
+ arr = con([True, False, True], dtype="boolean")
266
+ result = arr.to_numpy()
267
+ expected = np.array([True, False, True], dtype="bool")
268
+ tm.assert_numpy_array_equal(result, expected)
269
+
270
+ arr = con([True, False, None], dtype="boolean")
271
+ result = arr.to_numpy()
272
+ expected = np.array([True, False, pd.NA], dtype="object")
273
+ tm.assert_numpy_array_equal(result, expected)
274
+
275
+ arr = con([True, False, None], dtype="boolean")
276
+ result = arr.to_numpy(dtype="str")
277
+ expected = np.array([True, False, pd.NA], dtype=f"{tm.ENDIAN}U5")
278
+ tm.assert_numpy_array_equal(result, expected)
279
+
280
+ # no missing values -> can convert to bool, otherwise raises
281
+ arr = con([True, False, True], dtype="boolean")
282
+ result = arr.to_numpy(dtype="bool")
283
+ expected = np.array([True, False, True], dtype="bool")
284
+ tm.assert_numpy_array_equal(result, expected)
285
+
286
+ arr = con([True, False, None], dtype="boolean")
287
+ with pytest.raises(ValueError, match="cannot convert to 'bool'-dtype"):
288
+ result = arr.to_numpy(dtype="bool")
289
+
290
+ # specify dtype and na_value
291
+ arr = con([True, False, None], dtype="boolean")
292
+ result = arr.to_numpy(dtype=object, na_value=None)
293
+ expected = np.array([True, False, None], dtype="object")
294
+ tm.assert_numpy_array_equal(result, expected)
295
+
296
+ result = arr.to_numpy(dtype=bool, na_value=False)
297
+ expected = np.array([True, False, False], dtype="bool")
298
+ tm.assert_numpy_array_equal(result, expected)
299
+
300
+ result = arr.to_numpy(dtype="int64", na_value=-99)
301
+ expected = np.array([1, 0, -99], dtype="int64")
302
+ tm.assert_numpy_array_equal(result, expected)
303
+
304
+ result = arr.to_numpy(dtype="float64", na_value=np.nan)
305
+ expected = np.array([1, 0, np.nan], dtype="float64")
306
+ tm.assert_numpy_array_equal(result, expected)
307
+
308
+ # converting to int or float without specifying na_value raises
309
+ with pytest.raises(ValueError, match="cannot convert to 'int64'-dtype"):
310
+ arr.to_numpy(dtype="int64")
311
+
312
+
313
+ def test_to_numpy_copy():
314
+ # to_numpy can be zero-copy if no missing values
315
+ arr = pd.array([True, False, True], dtype="boolean")
316
+ result = arr.to_numpy(dtype=bool)
317
+ result[0] = False
318
+ tm.assert_extension_array_equal(
319
+ arr, pd.array([False, False, True], dtype="boolean")
320
+ )
321
+
322
+ arr = pd.array([True, False, True], dtype="boolean")
323
+ result = arr.to_numpy(dtype=bool, copy=True)
324
+ result[0] = False
325
+ tm.assert_extension_array_equal(arr, pd.array([True, False, True], dtype="boolean"))
326
+
327
+
328
+ def test_to_numpy_readonly():
329
+ arr = pd.array([True, False, True], dtype="boolean")
330
+ arr._readonly = True
331
+ result = arr.to_numpy(dtype=bool)
332
+ assert not result.flags.writeable
333
+
334
+ result = arr.to_numpy(dtype="int64")
335
+ assert result.flags.writeable
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_function.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+
7
+
8
+ @pytest.mark.parametrize(
9
+ "ufunc", [np.add, np.logical_or, np.logical_and, np.logical_xor]
10
+ )
11
+ def test_ufuncs_binary(ufunc):
12
+ # two BooleanArrays
13
+ a = pd.array([True, False, None], dtype="boolean")
14
+ result = ufunc(a, a)
15
+ expected = pd.array(ufunc(a._data, a._data), dtype="boolean")
16
+ expected[a._mask] = np.nan
17
+ tm.assert_extension_array_equal(result, expected)
18
+
19
+ s = pd.Series(a)
20
+ result = ufunc(s, a)
21
+ expected = pd.Series(ufunc(a._data, a._data), dtype="boolean")
22
+ expected[a._mask] = np.nan
23
+ tm.assert_series_equal(result, expected)
24
+
25
+ # Boolean with numpy array
26
+ arr = np.array([True, True, False])
27
+ result = ufunc(a, arr)
28
+ expected = pd.array(ufunc(a._data, arr), dtype="boolean")
29
+ expected[a._mask] = np.nan
30
+ tm.assert_extension_array_equal(result, expected)
31
+
32
+ result = ufunc(arr, a)
33
+ expected = pd.array(ufunc(arr, a._data), dtype="boolean")
34
+ expected[a._mask] = np.nan
35
+ tm.assert_extension_array_equal(result, expected)
36
+
37
+ # BooleanArray with scalar
38
+ result = ufunc(a, True)
39
+ expected = pd.array(ufunc(a._data, True), dtype="boolean")
40
+ expected[a._mask] = np.nan
41
+ tm.assert_extension_array_equal(result, expected)
42
+
43
+ result = ufunc(True, a)
44
+ expected = pd.array(ufunc(True, a._data), dtype="boolean")
45
+ expected[a._mask] = np.nan
46
+ tm.assert_extension_array_equal(result, expected)
47
+
48
+ # not handled types
49
+ msg = r"operand type\(s\) all returned NotImplemented from __array_ufunc__"
50
+ with pytest.raises(TypeError, match=msg):
51
+ ufunc(a, "test")
52
+
53
+
54
+ @pytest.mark.parametrize("ufunc", [np.logical_not])
55
+ def test_ufuncs_unary(ufunc):
56
+ a = pd.array([True, False, None], dtype="boolean")
57
+ result = ufunc(a)
58
+ expected = pd.array(ufunc(a._data), dtype="boolean")
59
+ expected[a._mask] = np.nan
60
+ tm.assert_extension_array_equal(result, expected)
61
+
62
+ ser = pd.Series(a)
63
+ result = ufunc(ser)
64
+ expected = pd.Series(ufunc(a._data), dtype="boolean")
65
+ expected[a._mask] = np.nan
66
+ tm.assert_series_equal(result, expected)
67
+
68
+
69
+ def test_ufunc_numeric():
70
+ # np.sqrt on np.bool_ returns float16, which we upcast to Float32
71
+ # bc we do not have Float16
72
+ arr = pd.array([True, False, None], dtype="boolean")
73
+
74
+ res = np.sqrt(arr)
75
+
76
+ expected = pd.array([1, 0, None], dtype="Float32")
77
+ tm.assert_extension_array_equal(res, expected)
78
+
79
+
80
+ @pytest.mark.parametrize("values", [[True, False], [True, None]])
81
+ def test_ufunc_reduce_raises(values):
82
+ arr = pd.array(values, dtype="boolean")
83
+
84
+ res = np.add.reduce(arr)
85
+ if arr[-1] is pd.NA:
86
+ expected = pd.NA
87
+ else:
88
+ expected = arr._data.sum()
89
+ tm.assert_almost_equal(res, expected)
90
+
91
+
92
+ def test_value_counts_na():
93
+ arr = pd.array([True, False, pd.NA], dtype="boolean")
94
+ result = arr.value_counts(dropna=False)
95
+ expected = pd.Series([1, 1, 1], index=arr, dtype="Int64", name="count")
96
+ assert expected.index.dtype == arr.dtype
97
+ tm.assert_series_equal(result, expected)
98
+
99
+ result = arr.value_counts(dropna=True)
100
+ expected = pd.Series([1, 1], index=arr[:-1], dtype="Int64", name="count")
101
+ assert expected.index.dtype == arr.dtype
102
+ tm.assert_series_equal(result, expected)
103
+
104
+
105
+ def test_value_counts_with_normalize():
106
+ ser = pd.Series([True, False, pd.NA], dtype="boolean")
107
+ result = ser.value_counts(normalize=True)
108
+ expected = pd.Series([1, 1], index=ser[:-1], dtype="Float64", name="proportion") / 2
109
+ assert expected.index.dtype == "boolean"
110
+ tm.assert_series_equal(result, expected)
111
+
112
+
113
+ def test_diff():
114
+ a = pd.array(
115
+ [True, True, False, False, True, None, True, None, False], dtype="boolean"
116
+ )
117
+ result = pd.core.algorithms.diff(a, 1)
118
+ expected = pd.array(
119
+ [None, False, True, False, True, None, None, None, None], dtype="boolean"
120
+ )
121
+ tm.assert_extension_array_equal(result, expected)
122
+
123
+ ser = pd.Series(a)
124
+ result = ser.diff()
125
+ expected = pd.Series(expected)
126
+ tm.assert_series_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_indexing.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+
7
+
8
+ @pytest.mark.parametrize("na", [None, np.nan, pd.NA])
9
+ def test_setitem_missing_values(na):
10
+ arr = pd.array([True, False, None], dtype="boolean")
11
+ expected = pd.array([True, None, None], dtype="boolean")
12
+ arr[1] = na
13
+ tm.assert_extension_array_equal(arr, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_logical.py ADDED
@@ -0,0 +1,255 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import operator
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ import pandas as pd
7
+ import pandas._testing as tm
8
+ from pandas.arrays import BooleanArray
9
+ from pandas.core.ops.mask_ops import (
10
+ kleene_and,
11
+ kleene_or,
12
+ kleene_xor,
13
+ )
14
+ from pandas.tests.extension.base import BaseOpsUtil
15
+
16
+
17
+ class TestLogicalOps(BaseOpsUtil):
18
+ def test_numpy_scalars_ok(self, all_logical_operators):
19
+ a = pd.array([True, False, None], dtype="boolean")
20
+ op = getattr(a, all_logical_operators)
21
+
22
+ tm.assert_extension_array_equal(op(True), op(np.bool_(True)))
23
+ tm.assert_extension_array_equal(op(False), op(np.bool_(False)))
24
+
25
+ def get_op_from_name(self, op_name):
26
+ short_opname = op_name.strip("_")
27
+ short_opname = short_opname if "xor" in short_opname else short_opname + "_"
28
+ try:
29
+ op = getattr(operator, short_opname)
30
+ except AttributeError:
31
+ # Assume it is the reverse operator
32
+ rop = getattr(operator, short_opname[1:])
33
+ op = lambda x, y: rop(y, x)
34
+
35
+ return op
36
+
37
+ def test_empty_ok(self, all_logical_operators):
38
+ a = pd.array([], dtype="boolean")
39
+ op_name = all_logical_operators
40
+ result = getattr(a, op_name)(True)
41
+ tm.assert_extension_array_equal(a, result)
42
+
43
+ result = getattr(a, op_name)(False)
44
+ tm.assert_extension_array_equal(a, result)
45
+
46
+ result = getattr(a, op_name)(pd.NA)
47
+ tm.assert_extension_array_equal(a, result)
48
+
49
+ @pytest.mark.parametrize(
50
+ "other", ["a", pd.Timestamp(2017, 1, 1, 12), np.timedelta64(4, "ns")]
51
+ )
52
+ def test_eq_mismatched_type(self, other):
53
+ # GH-44499
54
+ arr = pd.array([True, False])
55
+ result = arr == other
56
+ expected = pd.array([False, False])
57
+ tm.assert_extension_array_equal(result, expected)
58
+
59
+ result = arr != other
60
+ expected = pd.array([True, True])
61
+ tm.assert_extension_array_equal(result, expected)
62
+
63
+ @pytest.mark.parametrize("other", [[True, False], [True, False, True, False]])
64
+ def test_logical_length_mismatch_raises(self, other, all_logical_operators):
65
+ op_name = all_logical_operators
66
+ a = pd.array([True, False, None], dtype="boolean")
67
+ msg = "Lengths must match"
68
+
69
+ with pytest.raises(ValueError, match=msg):
70
+ getattr(a, op_name)(other)
71
+
72
+ with pytest.raises(ValueError, match=msg):
73
+ getattr(a, op_name)(np.array(other))
74
+
75
+ with pytest.raises(ValueError, match=msg):
76
+ getattr(a, op_name)(pd.array(other, dtype="boolean"))
77
+
78
+ def test_logical_nan_raises(self, all_logical_operators):
79
+ op_name = all_logical_operators
80
+ a = pd.array([True, False, None], dtype="boolean")
81
+ msg = "Got float instead"
82
+
83
+ with pytest.raises(TypeError, match=msg):
84
+ getattr(a, op_name)(np.nan)
85
+
86
+ @pytest.mark.parametrize("other", ["a", 1])
87
+ def test_non_bool_or_na_other_raises(self, other, all_logical_operators):
88
+ a = pd.array([True, False], dtype="boolean")
89
+ with pytest.raises(TypeError, match=str(type(other).__name__)):
90
+ getattr(a, all_logical_operators)(other)
91
+
92
+ def test_kleene_or(self):
93
+ # A clear test of behavior.
94
+ a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
95
+ b = pd.array([True, False, None] * 3, dtype="boolean")
96
+ result = a | b
97
+ expected = pd.array(
98
+ [True, True, True, True, False, None, True, None, None], dtype="boolean"
99
+ )
100
+ tm.assert_extension_array_equal(result, expected)
101
+
102
+ result = b | a
103
+ tm.assert_extension_array_equal(result, expected)
104
+
105
+ # ensure we haven't mutated anything inplace
106
+ tm.assert_extension_array_equal(
107
+ a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
108
+ )
109
+ tm.assert_extension_array_equal(
110
+ b, pd.array([True, False, None] * 3, dtype="boolean")
111
+ )
112
+
113
+ @pytest.mark.parametrize(
114
+ "other, expected",
115
+ [
116
+ (pd.NA, [True, None, None]),
117
+ (True, [True, True, True]),
118
+ (np.bool_(True), [True, True, True]),
119
+ (False, [True, False, None]),
120
+ (np.bool_(False), [True, False, None]),
121
+ ],
122
+ )
123
+ def test_kleene_or_scalar(self, other, expected):
124
+ # TODO: test True & False
125
+ a = pd.array([True, False, None], dtype="boolean")
126
+ result = a | other
127
+ expected = pd.array(expected, dtype="boolean")
128
+ tm.assert_extension_array_equal(result, expected)
129
+
130
+ result = other | a
131
+ tm.assert_extension_array_equal(result, expected)
132
+
133
+ # ensure we haven't mutated anything inplace
134
+ tm.assert_extension_array_equal(
135
+ a, pd.array([True, False, None], dtype="boolean")
136
+ )
137
+
138
+ def test_kleene_and(self):
139
+ # A clear test of behavior.
140
+ a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
141
+ b = pd.array([True, False, None] * 3, dtype="boolean")
142
+ result = a & b
143
+ expected = pd.array(
144
+ [True, False, None, False, False, False, None, False, None], dtype="boolean"
145
+ )
146
+ tm.assert_extension_array_equal(result, expected)
147
+
148
+ result = b & a
149
+ tm.assert_extension_array_equal(result, expected)
150
+
151
+ # ensure we haven't mutated anything inplace
152
+ tm.assert_extension_array_equal(
153
+ a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
154
+ )
155
+ tm.assert_extension_array_equal(
156
+ b, pd.array([True, False, None] * 3, dtype="boolean")
157
+ )
158
+
159
+ @pytest.mark.parametrize(
160
+ "other, expected",
161
+ [
162
+ (pd.NA, [None, False, None]),
163
+ (True, [True, False, None]),
164
+ (False, [False, False, False]),
165
+ (np.bool_(True), [True, False, None]),
166
+ (np.bool_(False), [False, False, False]),
167
+ ],
168
+ )
169
+ def test_kleene_and_scalar(self, other, expected):
170
+ a = pd.array([True, False, None], dtype="boolean")
171
+ result = a & other
172
+ expected = pd.array(expected, dtype="boolean")
173
+ tm.assert_extension_array_equal(result, expected)
174
+
175
+ result = other & a
176
+ tm.assert_extension_array_equal(result, expected)
177
+
178
+ # ensure we haven't mutated anything inplace
179
+ tm.assert_extension_array_equal(
180
+ a, pd.array([True, False, None], dtype="boolean")
181
+ )
182
+
183
+ def test_kleene_xor(self):
184
+ a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
185
+ b = pd.array([True, False, None] * 3, dtype="boolean")
186
+ result = a ^ b
187
+ expected = pd.array(
188
+ [False, True, None, True, False, None, None, None, None], dtype="boolean"
189
+ )
190
+ tm.assert_extension_array_equal(result, expected)
191
+
192
+ result = b ^ a
193
+ tm.assert_extension_array_equal(result, expected)
194
+
195
+ # ensure we haven't mutated anything inplace
196
+ tm.assert_extension_array_equal(
197
+ a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
198
+ )
199
+ tm.assert_extension_array_equal(
200
+ b, pd.array([True, False, None] * 3, dtype="boolean")
201
+ )
202
+
203
+ @pytest.mark.parametrize(
204
+ "other, expected",
205
+ [
206
+ (pd.NA, [None, None, None]),
207
+ (True, [False, True, None]),
208
+ (np.bool_(True), [False, True, None]),
209
+ (np.bool_(False), [True, False, None]),
210
+ ],
211
+ )
212
+ def test_kleene_xor_scalar(self, other, expected):
213
+ a = pd.array([True, False, None], dtype="boolean")
214
+ result = a ^ other
215
+ expected = pd.array(expected, dtype="boolean")
216
+ tm.assert_extension_array_equal(result, expected)
217
+
218
+ result = other ^ a
219
+ tm.assert_extension_array_equal(result, expected)
220
+
221
+ # ensure we haven't mutated anything inplace
222
+ tm.assert_extension_array_equal(
223
+ a, pd.array([True, False, None], dtype="boolean")
224
+ )
225
+
226
+ @pytest.mark.parametrize("other", [True, False, pd.NA, [True, False, None] * 3])
227
+ def test_no_masked_assumptions(self, other, all_logical_operators):
228
+ # The logical operations should not assume that masked values are False!
229
+ a = pd.arrays.BooleanArray(
230
+ np.array([True, True, True, False, False, False, True, False, True]),
231
+ np.array([False] * 6 + [True, True, True]),
232
+ )
233
+ b = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
234
+ if isinstance(other, list):
235
+ other = pd.array(other, dtype="boolean")
236
+
237
+ result = getattr(a, all_logical_operators)(other)
238
+ expected = getattr(b, all_logical_operators)(other)
239
+ tm.assert_extension_array_equal(result, expected)
240
+
241
+ if isinstance(other, BooleanArray):
242
+ other._data[other._mask] = True
243
+ a._data[a._mask] = False
244
+
245
+ result = getattr(a, all_logical_operators)(other)
246
+ expected = getattr(b, all_logical_operators)(other)
247
+ tm.assert_extension_array_equal(result, expected)
248
+
249
+
250
+ @pytest.mark.parametrize("operation", [kleene_or, kleene_xor, kleene_and])
251
+ def test_error_both_scalar(operation):
252
+ msg = r"Either `left` or `right` need to be an np\.ndarray."
253
+ with pytest.raises(TypeError, match=msg):
254
+ # masks need to be non-None, otherwise it ends up in an infinite recursion
255
+ operation(True, True, np.zeros(1), np.zeros(1))
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_ops.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+ import pandas._testing as tm
3
+
4
+
5
+ class TestUnaryOps:
6
+ def test_invert(self):
7
+ a = pd.array([True, False, None], dtype="boolean")
8
+ expected = pd.array([False, True, None], dtype="boolean")
9
+ tm.assert_extension_array_equal(~a, expected)
10
+
11
+ expected = pd.Series(expected, index=["a", "b", "c"], name="name")
12
+ result = ~pd.Series(a, index=["a", "b", "c"], name="name")
13
+ tm.assert_series_equal(result, expected)
14
+
15
+ df = pd.DataFrame({"A": a, "B": [True, False, False]}, index=["a", "b", "c"])
16
+ result = ~df
17
+ expected = pd.DataFrame(
18
+ {"A": expected, "B": [False, True, True]}, index=["a", "b", "c"]
19
+ )
20
+ tm.assert_frame_equal(result, expected)
21
+
22
+ def test_abs(self):
23
+ # matching numpy behavior, abs is the identity function
24
+ arr = pd.array([True, False, None], dtype="boolean")
25
+ result = abs(arr)
26
+
27
+ tm.assert_extension_array_equal(result, arr)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_reduction.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+
6
+
7
+ @pytest.fixture
8
+ def data():
9
+ """Fixture returning boolean array, with valid and missing values."""
10
+ return pd.array(
11
+ [True, False] * 2 + [np.nan] + [True, False] + [np.nan] + [True, False],
12
+ dtype="boolean",
13
+ )
14
+
15
+
16
+ @pytest.mark.parametrize(
17
+ "values, exp_any, exp_all, exp_any_noskip, exp_all_noskip",
18
+ [
19
+ ([True, pd.NA], True, True, True, pd.NA),
20
+ ([False, pd.NA], False, False, pd.NA, False),
21
+ ([pd.NA], False, True, pd.NA, pd.NA),
22
+ ([], False, True, False, True),
23
+ # GH-33253: all True / all False values buggy with skipna=False
24
+ ([True, True], True, True, True, True),
25
+ ([False, False], False, False, False, False),
26
+ ],
27
+ )
28
+ @pytest.mark.parametrize("con", [pd.array, pd.Series])
29
+ def test_any_all(
30
+ values, exp_any, exp_all, exp_any_noskip, exp_all_noskip, using_python_scalars, con
31
+ ):
32
+ # the methods return numpy scalars
33
+ if not using_python_scalars or con is pd.array:
34
+ exp_any = pd.NA if exp_any is pd.NA else np.bool_(exp_any)
35
+ exp_all = pd.NA if exp_all is pd.NA else np.bool_(exp_all)
36
+ exp_any_noskip = pd.NA if exp_any_noskip is pd.NA else np.bool_(exp_any_noskip)
37
+ exp_all_noskip = pd.NA if exp_all_noskip is pd.NA else np.bool_(exp_all_noskip)
38
+
39
+ a = con(values, dtype="boolean")
40
+ assert a.any() is exp_any
41
+ assert a.all() is exp_all
42
+ assert a.any(skipna=False) is exp_any_noskip
43
+ assert a.all(skipna=False) is exp_all_noskip
44
+
45
+
46
+ def test_reductions_return_types(
47
+ dropna, data, all_numeric_reductions, using_python_scalars
48
+ ):
49
+ op = all_numeric_reductions
50
+ s = pd.Series(data)
51
+ if dropna:
52
+ s = s.dropna()
53
+
54
+ if using_python_scalars:
55
+ expected = {
56
+ "sum": int,
57
+ "prod": int,
58
+ "count": int,
59
+ "min": bool,
60
+ "max": bool,
61
+ }.get(op, float)
62
+ else:
63
+ expected = {
64
+ "sum": np.int_,
65
+ "prod": np.int_,
66
+ "count": np.integer,
67
+ "min": np.bool_,
68
+ "max": np.bool_,
69
+ }.get(op, np.float64)
70
+ result = getattr(s, op)()
71
+ assert isinstance(result, expected), f"{type(result)} vs {expected}"
.venv/lib/python3.13/site-packages/pandas/tests/arrays/boolean/test_repr.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+
3
+
4
+ def test_repr():
5
+ df = pd.DataFrame({"A": pd.array([True, False, None], dtype="boolean")})
6
+ expected = " A\n0 True\n1 False\n2 <NA>"
7
+ assert repr(df) == expected
8
+
9
+ expected = "0 True\n1 False\n2 <NA>\nName: A, dtype: boolean"
10
+ assert repr(df.A) == expected
11
+
12
+ expected = "<BooleanArray>\n[True, False, <NA>]\nLength: 3, dtype: boolean"
13
+ assert repr(df.A.array) == expected
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/__init__.py ADDED
File without changes
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_algos.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ import pandas._testing as tm
6
+
7
+
8
+ @pytest.mark.parametrize("categories", [["b", "a", "c"], ["a", "b", "c", "d"]])
9
+ def test_factorize(categories, ordered):
10
+ cat = pd.Categorical(
11
+ ["b", "b", "a", "c", None], categories=categories, ordered=ordered
12
+ )
13
+ codes, uniques = pd.factorize(cat)
14
+ expected_codes = np.array([0, 0, 1, 2, -1], dtype=np.intp)
15
+ expected_uniques = pd.Categorical(
16
+ ["b", "a", "c"], categories=categories, ordered=ordered
17
+ )
18
+
19
+ tm.assert_numpy_array_equal(codes, expected_codes)
20
+ tm.assert_categorical_equal(uniques, expected_uniques)
21
+
22
+
23
+ def test_factorized_sort():
24
+ cat = pd.Categorical(["b", "b", None, "a"])
25
+ codes, uniques = pd.factorize(cat, sort=True)
26
+ expected_codes = np.array([1, 1, -1, 0], dtype=np.intp)
27
+ expected_uniques = pd.Categorical(["a", "b"])
28
+
29
+ tm.assert_numpy_array_equal(codes, expected_codes)
30
+ tm.assert_categorical_equal(uniques, expected_uniques)
31
+
32
+
33
+ def test_factorized_sort_ordered():
34
+ cat = pd.Categorical(
35
+ ["b", "b", None, "a"], categories=["c", "b", "a"], ordered=True
36
+ )
37
+
38
+ codes, uniques = pd.factorize(cat, sort=True)
39
+ expected_codes = np.array([0, 0, -1, 1], dtype=np.intp)
40
+ expected_uniques = pd.Categorical(
41
+ ["b", "a"], categories=["c", "b", "a"], ordered=True
42
+ )
43
+
44
+ tm.assert_numpy_array_equal(codes, expected_codes)
45
+ tm.assert_categorical_equal(uniques, expected_uniques)
46
+
47
+
48
+ def test_isin_cats():
49
+ # GH2003
50
+ cat = pd.Categorical(["a", "b", np.nan])
51
+
52
+ result = cat.isin(["a", np.nan])
53
+ expected = np.array([True, False, True], dtype=bool)
54
+ tm.assert_numpy_array_equal(expected, result)
55
+
56
+ result = cat.isin(["a", "c"])
57
+ expected = np.array([True, False, False], dtype=bool)
58
+ tm.assert_numpy_array_equal(expected, result)
59
+
60
+
61
+ @pytest.mark.parametrize("value", [[""], [None, ""], [pd.NaT, ""]])
62
+ def test_isin_cats_corner_cases(value):
63
+ # GH36550
64
+ cat = pd.Categorical([""])
65
+ result = cat.isin(value)
66
+ expected = np.array([True], dtype=bool)
67
+ tm.assert_numpy_array_equal(expected, result)
68
+
69
+
70
+ @pytest.mark.parametrize("empty", [[], pd.Series(dtype=object), np.array([])])
71
+ def test_isin_empty(empty):
72
+ s = pd.Categorical(["a", "b"])
73
+ expected = np.array([False, False], dtype=bool)
74
+
75
+ result = s.isin(empty)
76
+ tm.assert_numpy_array_equal(expected, result)
77
+
78
+
79
+ def test_diff():
80
+ ser = pd.Series([1, 2, 3], dtype="category")
81
+
82
+ msg = "Convert to a suitable dtype"
83
+ with pytest.raises(TypeError, match=msg):
84
+ ser.diff()
85
+
86
+ df = ser.to_frame(name="A")
87
+ with pytest.raises(TypeError, match=msg):
88
+ df.diff()
89
+
90
+
91
+ def test_hash_read_only_categorical():
92
+ # GH#58481
93
+ idx = pd.Index(pd.Index(["a", "b", "c"], dtype="object").values)
94
+ cat = pd.CategoricalDtype(idx)
95
+ arr = pd.Series(["a", "b"], dtype=cat).values
96
+ assert hash(arr.dtype) == hash(arr.dtype)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_analytics.py ADDED
@@ -0,0 +1,353 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import sys
3
+
4
+ import numpy as np
5
+ import pytest
6
+
7
+ from pandas.compat import PYPY
8
+
9
+ from pandas import (
10
+ Categorical,
11
+ CategoricalDtype,
12
+ DataFrame,
13
+ Index,
14
+ NaT,
15
+ Series,
16
+ date_range,
17
+ )
18
+ import pandas._testing as tm
19
+ from pandas.api.types import is_scalar
20
+
21
+
22
+ class TestCategoricalAnalytics:
23
+ @pytest.mark.parametrize("aggregation", ["min", "max"])
24
+ def test_min_max_not_ordered_raises(self, aggregation):
25
+ # unordered cats have no min/max
26
+ cat = Categorical(["a", "b", "c", "d"], ordered=False)
27
+ msg = f"Categorical is not ordered for operation {aggregation}"
28
+ agg_func = getattr(cat, aggregation)
29
+
30
+ with pytest.raises(TypeError, match=msg):
31
+ agg_func()
32
+
33
+ ufunc = np.minimum if aggregation == "min" else np.maximum
34
+ with pytest.raises(TypeError, match=msg):
35
+ ufunc.reduce(cat)
36
+
37
+ def test_min_max_ordered(self, index_or_series_or_array):
38
+ cat = Categorical(["a", "b", "c", "d"], ordered=True)
39
+ obj = index_or_series_or_array(cat)
40
+ _min = obj.min()
41
+ _max = obj.max()
42
+ assert _min == "a"
43
+ assert _max == "d"
44
+
45
+ assert np.minimum.reduce(obj) == "a"
46
+ assert np.maximum.reduce(obj) == "d"
47
+ # TODO: raises if we pass axis=0 (on Index and Categorical, not Series)
48
+
49
+ cat = Categorical(
50
+ ["a", "b", "c", "d"], categories=["d", "c", "b", "a"], ordered=True
51
+ )
52
+ obj = index_or_series_or_array(cat)
53
+ _min = obj.min()
54
+ _max = obj.max()
55
+ assert _min == "d"
56
+ assert _max == "a"
57
+ assert np.minimum.reduce(obj) == "d"
58
+ assert np.maximum.reduce(obj) == "a"
59
+
60
+ def test_min_max_reduce(self):
61
+ # GH52788
62
+ cat = Categorical(["a", "b", "c", "d"], ordered=True)
63
+ df = DataFrame(cat)
64
+
65
+ result_max = df.agg("max")
66
+ expected_max = Series(Categorical(["d"], dtype=cat.dtype))
67
+ tm.assert_series_equal(result_max, expected_max)
68
+
69
+ result_min = df.agg("min")
70
+ expected_min = Series(Categorical(["a"], dtype=cat.dtype))
71
+ tm.assert_series_equal(result_min, expected_min)
72
+
73
+ @pytest.mark.parametrize(
74
+ "categories,expected",
75
+ [
76
+ (list("ABC"), np.nan),
77
+ ([1, 2, 3], np.nan),
78
+ pytest.param(
79
+ Series(date_range("2020-01-01", periods=3), dtype="category"),
80
+ NaT,
81
+ marks=pytest.mark.xfail(
82
+ reason="https://github.com/pandas-dev/pandas/issues/29962"
83
+ ),
84
+ ),
85
+ ],
86
+ )
87
+ @pytest.mark.parametrize("aggregation", ["min", "max"])
88
+ def test_min_max_ordered_empty(self, categories, expected, aggregation):
89
+ # GH 30227
90
+ cat = Categorical([], categories=categories, ordered=True)
91
+
92
+ agg_func = getattr(cat, aggregation)
93
+ result = agg_func()
94
+ assert result is expected
95
+
96
+ @pytest.mark.parametrize(
97
+ "values, categories",
98
+ [(["a", "b", "c", np.nan], list("cba")), ([1, 2, 3, np.nan], [3, 2, 1])],
99
+ )
100
+ @pytest.mark.parametrize("function", ["min", "max"])
101
+ def test_min_max_with_nan(self, values, categories, function, skipna):
102
+ # GH 25303
103
+ cat = Categorical(values, categories=categories, ordered=True)
104
+ result = getattr(cat, function)(skipna=skipna)
105
+
106
+ if skipna is False:
107
+ assert result is np.nan
108
+ else:
109
+ expected = categories[0] if function == "min" else categories[2]
110
+ assert result == expected
111
+
112
+ @pytest.mark.parametrize("function", ["min", "max"])
113
+ def test_min_max_only_nan(self, function, skipna):
114
+ # https://github.com/pandas-dev/pandas/issues/33450
115
+ cat = Categorical([np.nan], categories=[1, 2], ordered=True)
116
+ result = getattr(cat, function)(skipna=skipna)
117
+ assert result is np.nan
118
+
119
+ @pytest.mark.parametrize("method", ["min", "max"])
120
+ def test_numeric_only_min_max_raises(self, method):
121
+ # GH 25303
122
+ cat = Categorical(
123
+ [np.nan, 1, 2, np.nan], categories=[5, 4, 3, 2, 1], ordered=True
124
+ )
125
+ with pytest.raises(TypeError, match=".* got an unexpected keyword"):
126
+ getattr(cat, method)(numeric_only=True)
127
+
128
+ @pytest.mark.parametrize("method", ["min", "max"])
129
+ def test_numpy_min_max_raises(self, method):
130
+ cat = Categorical(["a", "b", "c", "b"], ordered=False)
131
+ msg = (
132
+ f"Categorical is not ordered for operation {method}\n"
133
+ "you can use .as_ordered() to change the Categorical to an ordered one"
134
+ )
135
+ method = getattr(np, method)
136
+ with pytest.raises(TypeError, match=re.escape(msg)):
137
+ method(cat)
138
+
139
+ @pytest.mark.parametrize("kwarg", ["axis", "out", "keepdims"])
140
+ @pytest.mark.parametrize("method", ["min", "max"])
141
+ def test_numpy_min_max_unsupported_kwargs_raises(self, method, kwarg):
142
+ cat = Categorical(["a", "b", "c", "b"], ordered=True)
143
+ msg = (
144
+ f"the '{kwarg}' parameter is not supported in the pandas implementation "
145
+ f"of {method}"
146
+ )
147
+ if kwarg == "axis":
148
+ msg = r"`axis` must be fewer than the number of dimensions \(1\)"
149
+ kwargs = {kwarg: 42}
150
+ method = getattr(np, method)
151
+ with pytest.raises(ValueError, match=msg):
152
+ method(cat, **kwargs)
153
+
154
+ @pytest.mark.parametrize("method, expected", [("min", "a"), ("max", "c")])
155
+ def test_numpy_min_max_axis_equals_none(self, method, expected):
156
+ cat = Categorical(["a", "b", "c", "b"], ordered=True)
157
+ method = getattr(np, method)
158
+ result = method(cat, axis=None)
159
+ assert result == expected
160
+
161
+ @pytest.mark.parametrize(
162
+ "values,categories,exp_mode",
163
+ [
164
+ ([1, 1, 2, 4, 5, 5, 5], [5, 4, 3, 2, 1], [5]),
165
+ ([1, 1, 1, 4, 5, 5, 5], [5, 4, 3, 2, 1], [5, 1]),
166
+ ([1, 2, 3, 4, 5], [5, 4, 3, 2, 1], [5, 4, 3, 2, 1]),
167
+ ([np.nan, np.nan, np.nan, 4, 5], [5, 4, 3, 2, 1], [5, 4]),
168
+ ([np.nan, np.nan, np.nan, 4, 5, 4], [5, 4, 3, 2, 1], [4]),
169
+ ([np.nan, np.nan, 4, 5, 4], [5, 4, 3, 2, 1], [4]),
170
+ ],
171
+ )
172
+ def test_mode(self, values, categories, exp_mode):
173
+ cat = Categorical(values, categories=categories, ordered=True)
174
+ res = Series(cat).mode()._values
175
+ exp = Categorical(exp_mode, categories=categories, ordered=True)
176
+ tm.assert_categorical_equal(res, exp)
177
+
178
+ def test_searchsorted(self, ordered):
179
+ # https://github.com/pandas-dev/pandas/issues/8420
180
+ # https://github.com/pandas-dev/pandas/issues/14522
181
+
182
+ cat = Categorical(
183
+ ["cheese", "milk", "apple", "bread", "bread"],
184
+ categories=["cheese", "milk", "apple", "bread"],
185
+ ordered=ordered,
186
+ )
187
+ ser = Series(cat)
188
+
189
+ # Searching for single item argument, side='left' (default)
190
+ res_cat = cat.searchsorted("apple")
191
+ assert res_cat == 2
192
+ assert is_scalar(res_cat)
193
+
194
+ res_ser = ser.searchsorted("apple")
195
+ assert res_ser == 2
196
+ assert is_scalar(res_ser)
197
+
198
+ # Searching for single item array, side='left' (default)
199
+ res_cat = cat.searchsorted(["bread"])
200
+ res_ser = ser.searchsorted(["bread"])
201
+ exp = np.array([3], dtype=np.intp)
202
+ tm.assert_numpy_array_equal(res_cat, exp)
203
+ tm.assert_numpy_array_equal(res_ser, exp)
204
+
205
+ # Searching for several items array, side='right'
206
+ res_cat = cat.searchsorted(["apple", "bread"], side="right")
207
+ res_ser = ser.searchsorted(["apple", "bread"], side="right")
208
+ exp = np.array([3, 5], dtype=np.intp)
209
+ tm.assert_numpy_array_equal(res_cat, exp)
210
+ tm.assert_numpy_array_equal(res_ser, exp)
211
+
212
+ # Searching for a single value that is not from the Categorical
213
+ with pytest.raises(TypeError, match="cucumber"):
214
+ cat.searchsorted("cucumber")
215
+ with pytest.raises(TypeError, match="cucumber"):
216
+ ser.searchsorted("cucumber")
217
+
218
+ # Searching for multiple values one of each is not from the Categorical
219
+ msg = (
220
+ "Cannot setitem on a Categorical with a new category, "
221
+ "set the categories first"
222
+ )
223
+ with pytest.raises(TypeError, match=msg):
224
+ cat.searchsorted(["bread", "cucumber"])
225
+ with pytest.raises(TypeError, match=msg):
226
+ ser.searchsorted(["bread", "cucumber"])
227
+
228
+ def test_unique(self, ordered):
229
+ # GH38140
230
+ dtype = CategoricalDtype(["a", "b", "c"], ordered=ordered)
231
+
232
+ # categories are reordered based on value when ordered=False
233
+ cat = Categorical(["a", "b", "c"], dtype=dtype)
234
+ res = cat.unique()
235
+ tm.assert_categorical_equal(res, cat)
236
+
237
+ cat = Categorical(["a", "b", "a", "a"], dtype=dtype)
238
+ res = cat.unique()
239
+ tm.assert_categorical_equal(res, Categorical(["a", "b"], dtype=dtype))
240
+
241
+ cat = Categorical(["c", "a", "b", "a", "a"], dtype=dtype)
242
+ res = cat.unique()
243
+ exp_cat = Categorical(["c", "a", "b"], dtype=dtype)
244
+ tm.assert_categorical_equal(res, exp_cat)
245
+
246
+ # nan must be removed
247
+ cat = Categorical(["b", np.nan, "b", np.nan, "a"], dtype=dtype)
248
+ res = cat.unique()
249
+ exp_cat = Categorical(["b", np.nan, "a"], dtype=dtype)
250
+ tm.assert_categorical_equal(res, exp_cat)
251
+
252
+ def test_unique_index_series(self, ordered):
253
+ # GH38140
254
+ dtype = CategoricalDtype([3, 2, 1], ordered=ordered)
255
+
256
+ c = Categorical([3, 1, 2, 2, 1], dtype=dtype)
257
+ # Categorical.unique sorts categories by appearance order
258
+ # if ordered=False
259
+ exp = Categorical([3, 1, 2], dtype=dtype)
260
+ tm.assert_categorical_equal(c.unique(), exp)
261
+
262
+ tm.assert_index_equal(Index(c).unique(), Index(exp))
263
+ tm.assert_categorical_equal(Series(c).unique(), exp)
264
+
265
+ c = Categorical([1, 1, 2, 2], dtype=dtype)
266
+ exp = Categorical([1, 2], dtype=dtype)
267
+ tm.assert_categorical_equal(c.unique(), exp)
268
+ tm.assert_index_equal(Index(c).unique(), Index(exp))
269
+ tm.assert_categorical_equal(Series(c).unique(), exp)
270
+
271
+ def test_shift(self):
272
+ # GH 9416
273
+ cat = Categorical(["a", "b", "c", "d", "a"])
274
+
275
+ # shift forward
276
+ sp1 = cat.shift(1)
277
+ xp1 = Categorical([np.nan, "a", "b", "c", "d"])
278
+ tm.assert_categorical_equal(sp1, xp1)
279
+ tm.assert_categorical_equal(cat[:-1], sp1[1:])
280
+
281
+ # shift back
282
+ sn2 = cat.shift(-2)
283
+ xp2 = Categorical(
284
+ ["c", "d", "a", np.nan, np.nan], categories=["a", "b", "c", "d"]
285
+ )
286
+ tm.assert_categorical_equal(sn2, xp2)
287
+ tm.assert_categorical_equal(cat[2:], sn2[:-2])
288
+
289
+ # shift by zero
290
+ tm.assert_categorical_equal(cat, cat.shift(0))
291
+
292
+ def test_nbytes(self):
293
+ cat = Categorical([1, 2, 3])
294
+ exp = 3 + 3 * 8 # 3 int8s for values + 3 int64s for categories
295
+ assert cat.nbytes == exp
296
+
297
+ def test_memory_usage(self, using_infer_string):
298
+ cat = Categorical([1, 2, 3])
299
+
300
+ # .categories is an index, so we include the hashtable
301
+ assert 0 < cat.nbytes <= cat.memory_usage()
302
+ assert 0 < cat.nbytes <= cat.memory_usage(deep=True)
303
+
304
+ cat = Categorical(["foo", "foo", "bar"])
305
+ if using_infer_string:
306
+ if cat.categories.dtype.storage == "python":
307
+ assert cat.memory_usage(deep=True) > cat.nbytes
308
+ else:
309
+ assert cat.memory_usage(deep=True) >= cat.nbytes
310
+ else:
311
+ assert cat.memory_usage(deep=True) > cat.nbytes
312
+
313
+ if not PYPY:
314
+ # sys.getsizeof will call the .memory_usage with
315
+ # deep=True, and add on some GC overhead
316
+ diff = cat.memory_usage(deep=True) - sys.getsizeof(cat)
317
+ assert abs(diff) < 100
318
+
319
+ def test_map(self):
320
+ c = Categorical(list("ABABC"), categories=list("CBA"), ordered=True)
321
+ result = c.map(lambda x: x.lower(), na_action=None)
322
+ exp = Categorical(list("ababc"), categories=list("cba"), ordered=True)
323
+ tm.assert_categorical_equal(result, exp)
324
+
325
+ c = Categorical(list("ABABC"), categories=list("ABC"), ordered=False)
326
+ result = c.map(lambda x: x.lower(), na_action=None)
327
+ exp = Categorical(list("ababc"), categories=list("abc"), ordered=False)
328
+ tm.assert_categorical_equal(result, exp)
329
+
330
+ result = c.map(lambda x: 1, na_action=None)
331
+ # GH 12766: Return an index not an array
332
+ tm.assert_index_equal(result, Index(np.array([1] * 5, dtype=np.int64)))
333
+
334
+ @pytest.mark.parametrize("value", [1, "True", [1, 2, 3], 5.0])
335
+ def test_validate_inplace_raises(self, value):
336
+ cat = Categorical(["A", "B", "B", "C", "A"])
337
+ msg = (
338
+ 'For argument "inplace" expected type bool, '
339
+ f"received type {type(value).__name__}"
340
+ )
341
+
342
+ with pytest.raises(ValueError, match=msg):
343
+ cat.sort_values(inplace=value)
344
+
345
+ def test_quantile_empty(self):
346
+ # make sure we have correct itemsize on resulting codes
347
+ cat = Categorical(["A", "B"])
348
+ idx = Index([0.0, 0.5])
349
+ result = cat[:0]._quantile(idx, interpolation="linear")
350
+ assert result._codes.dtype == np.int8
351
+
352
+ expected = cat.take([-1, -1], allow_fill=True)
353
+ tm.assert_extension_array_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_api.py ADDED
@@ -0,0 +1,493 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ from pandas.errors import Pandas4Warning
7
+
8
+ from pandas import (
9
+ Categorical,
10
+ CategoricalIndex,
11
+ DataFrame,
12
+ Index,
13
+ Series,
14
+ StringDtype,
15
+ )
16
+ import pandas._testing as tm
17
+ from pandas.core.arrays.categorical import recode_for_categories
18
+
19
+
20
+ class TestCategoricalAPI:
21
+ def test_ordered_api(self):
22
+ # GH 9347
23
+ cat1 = Categorical(list("acb"), ordered=False)
24
+ tm.assert_index_equal(cat1.categories, Index(["a", "b", "c"]))
25
+ assert not cat1.ordered
26
+
27
+ cat2 = Categorical(list("acb"), categories=list("bca"), ordered=False)
28
+ tm.assert_index_equal(cat2.categories, Index(["b", "c", "a"]))
29
+ assert not cat2.ordered
30
+
31
+ cat3 = Categorical(list("acb"), ordered=True)
32
+ tm.assert_index_equal(cat3.categories, Index(["a", "b", "c"]))
33
+ assert cat3.ordered
34
+
35
+ cat4 = Categorical(list("acb"), categories=list("bca"), ordered=True)
36
+ tm.assert_index_equal(cat4.categories, Index(["b", "c", "a"]))
37
+ assert cat4.ordered
38
+
39
+ def test_set_ordered(self):
40
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
41
+ cat2 = cat.as_unordered()
42
+ assert not cat2.ordered
43
+ cat2 = cat.as_ordered()
44
+ assert cat2.ordered
45
+
46
+ assert cat2.set_ordered(True).ordered
47
+ assert not cat2.set_ordered(False).ordered
48
+
49
+ # removed in 0.19.0
50
+ msg = "property 'ordered' of 'Categorical' object has no setter"
51
+ with pytest.raises(AttributeError, match=msg):
52
+ cat.ordered = True
53
+ with pytest.raises(AttributeError, match=msg):
54
+ cat.ordered = False
55
+
56
+ def test_rename_categories(self):
57
+ cat = Categorical(["a", "b", "c", "a"])
58
+
59
+ # inplace=False: the old one must not be changed
60
+ res = cat.rename_categories([1, 2, 3])
61
+ tm.assert_numpy_array_equal(
62
+ res.__array__(), np.array([1, 2, 3, 1], dtype=np.int64)
63
+ )
64
+ tm.assert_index_equal(res.categories, Index([1, 2, 3]))
65
+
66
+ exp_cat = np.array(["a", "b", "c", "a"], dtype=np.object_)
67
+ tm.assert_numpy_array_equal(cat.__array__(), exp_cat)
68
+
69
+ exp_cat = Index(["a", "b", "c"])
70
+ tm.assert_index_equal(cat.categories, exp_cat)
71
+
72
+ # GH18862 (let rename_categories take callables)
73
+ result = cat.rename_categories(lambda x: x.upper())
74
+ expected = Categorical(["A", "B", "C", "A"])
75
+ tm.assert_categorical_equal(result, expected)
76
+
77
+ @pytest.mark.parametrize("new_categories", [[1, 2, 3, 4], [1, 2]])
78
+ def test_rename_categories_wrong_length_raises(self, new_categories):
79
+ cat = Categorical(["a", "b", "c", "a"])
80
+ msg = (
81
+ "new categories need to have the same number of items as the "
82
+ "old categories!"
83
+ )
84
+ with pytest.raises(ValueError, match=msg):
85
+ cat.rename_categories(new_categories)
86
+
87
+ def test_rename_categories_series(self):
88
+ # https://github.com/pandas-dev/pandas/issues/17981
89
+ c = Categorical(["a", "b"])
90
+ result = c.rename_categories(Series([0, 1], index=["a", "b"]))
91
+ expected = Categorical([0, 1])
92
+ tm.assert_categorical_equal(result, expected)
93
+
94
+ def test_rename_categories_dict(self):
95
+ # GH 17336
96
+ cat = Categorical(["a", "b", "c", "d"])
97
+ res = cat.rename_categories({"a": 4, "b": 3, "c": 2, "d": 1})
98
+ expected = Index([4, 3, 2, 1])
99
+ tm.assert_index_equal(res.categories, expected)
100
+
101
+ # Test for dicts of smaller length
102
+ cat = Categorical(["a", "b", "c", "d"])
103
+ res = cat.rename_categories({"a": 1, "c": 3})
104
+
105
+ expected = Index([1, "b", 3, "d"])
106
+ tm.assert_index_equal(res.categories, expected)
107
+
108
+ # Test for dicts with bigger length
109
+ cat = Categorical(["a", "b", "c", "d"])
110
+ res = cat.rename_categories({"a": 1, "b": 2, "c": 3, "d": 4, "e": 5, "f": 6})
111
+ expected = Index([1, 2, 3, 4])
112
+ tm.assert_index_equal(res.categories, expected)
113
+
114
+ # Test for dicts with no items from old categories
115
+ cat = Categorical(["a", "b", "c", "d"])
116
+ res = cat.rename_categories({"f": 1, "g": 3})
117
+
118
+ expected = Index(["a", "b", "c", "d"])
119
+ tm.assert_index_equal(res.categories, expected)
120
+
121
+ def test_reorder_categories(self):
122
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
123
+ old = cat.copy()
124
+ new = Categorical(
125
+ ["a", "b", "c", "a"], categories=["c", "b", "a"], ordered=True
126
+ )
127
+
128
+ res = cat.reorder_categories(["c", "b", "a"])
129
+ # cat must be the same as before
130
+ tm.assert_categorical_equal(cat, old)
131
+ # only res is changed
132
+ tm.assert_categorical_equal(res, new)
133
+
134
+ @pytest.mark.parametrize(
135
+ "new_categories",
136
+ [
137
+ ["a"], # not all "old" included in "new"
138
+ ["a", "b", "d"], # still not all "old" in "new"
139
+ ["a", "b", "c", "d"], # all "old" included in "new", but too long
140
+ ],
141
+ )
142
+ def test_reorder_categories_raises(self, new_categories):
143
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
144
+ msg = "items in new_categories are not the same as in old categories"
145
+ with pytest.raises(ValueError, match=msg):
146
+ cat.reorder_categories(new_categories)
147
+
148
+ def test_add_categories(self):
149
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
150
+ old = cat.copy()
151
+ new = Categorical(
152
+ ["a", "b", "c", "a"], categories=["a", "b", "c", "d"], ordered=True
153
+ )
154
+
155
+ res = cat.add_categories("d")
156
+ tm.assert_categorical_equal(cat, old)
157
+ tm.assert_categorical_equal(res, new)
158
+
159
+ res = cat.add_categories(["d"])
160
+ tm.assert_categorical_equal(cat, old)
161
+ tm.assert_categorical_equal(res, new)
162
+
163
+ # GH 9927
164
+ cat = Categorical(list("abc"), ordered=True)
165
+ expected = Categorical(list("abc"), categories=list("abcde"), ordered=True)
166
+ # test with Series, np.array, index, list
167
+ res = cat.add_categories(Series(["d", "e"]))
168
+ tm.assert_categorical_equal(res, expected)
169
+ res = cat.add_categories(np.array(["d", "e"]))
170
+ tm.assert_categorical_equal(res, expected)
171
+ res = cat.add_categories(Index(["d", "e"]))
172
+ tm.assert_categorical_equal(res, expected)
173
+ res = cat.add_categories(["d", "e"])
174
+ tm.assert_categorical_equal(res, expected)
175
+
176
+ def test_add_categories_existing_raises(self):
177
+ # new is in old categories
178
+ cat = Categorical(["a", "b", "c", "d"], ordered=True)
179
+ msg = re.escape("new categories must not include old categories: {'d'}")
180
+ with pytest.raises(ValueError, match=msg):
181
+ cat.add_categories(["d"])
182
+
183
+ def test_add_categories_losing_dtype_information(self):
184
+ # GH#48812
185
+ cat = Categorical(Series([1, 2], dtype="Int64"))
186
+ ser = Series([4], dtype="Int64")
187
+ result = cat.add_categories(ser)
188
+ expected = Categorical(
189
+ Series([1, 2], dtype="Int64"), categories=Series([1, 2, 4], dtype="Int64")
190
+ )
191
+ tm.assert_categorical_equal(result, expected)
192
+
193
+ cat = Categorical(Series(["a", "b", "a"], dtype=StringDtype()))
194
+ ser = Series(["d"], dtype=StringDtype())
195
+ result = cat.add_categories(ser)
196
+ expected = Categorical(
197
+ Series(["a", "b", "a"], dtype=StringDtype()),
198
+ categories=Series(["a", "b", "d"], dtype=StringDtype()),
199
+ )
200
+ tm.assert_categorical_equal(result, expected)
201
+
202
+ def test_set_categories(self):
203
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
204
+ exp_categories = Index(["c", "b", "a"])
205
+ exp_values = np.array(["a", "b", "c", "a"], dtype=np.object_)
206
+
207
+ cat = cat.set_categories(["c", "b", "a"])
208
+ res = cat.set_categories(["a", "b", "c"])
209
+ # cat must be the same as before
210
+ tm.assert_index_equal(cat.categories, exp_categories)
211
+ tm.assert_numpy_array_equal(cat.__array__(), exp_values)
212
+ # only res is changed
213
+ exp_categories_back = Index(["a", "b", "c"])
214
+ tm.assert_index_equal(res.categories, exp_categories_back)
215
+ tm.assert_numpy_array_equal(res.__array__(), exp_values)
216
+
217
+ # not all "old" included in "new" -> all not included ones are now
218
+ # np.nan
219
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
220
+ res = cat.set_categories(["a"])
221
+ tm.assert_numpy_array_equal(res.codes, np.array([0, -1, -1, 0], dtype=np.int8))
222
+
223
+ # still not all "old" in "new"
224
+ res = cat.set_categories(["a", "b", "d"])
225
+ tm.assert_numpy_array_equal(res.codes, np.array([0, 1, -1, 0], dtype=np.int8))
226
+ tm.assert_index_equal(res.categories, Index(["a", "b", "d"]))
227
+
228
+ # all "old" included in "new"
229
+ cat = cat.set_categories(["a", "b", "c", "d"])
230
+ exp_categories = Index(["a", "b", "c", "d"])
231
+ tm.assert_index_equal(cat.categories, exp_categories)
232
+
233
+ # internals...
234
+ c = Categorical([1, 2, 3, 4, 1], categories=[1, 2, 3, 4], ordered=True)
235
+ tm.assert_numpy_array_equal(c._codes, np.array([0, 1, 2, 3, 0], dtype=np.int8))
236
+ tm.assert_index_equal(c.categories, Index([1, 2, 3, 4]))
237
+
238
+ exp = np.array([1, 2, 3, 4, 1], dtype=np.int64)
239
+ tm.assert_numpy_array_equal(np.asarray(c), exp)
240
+
241
+ # all "pointers" to '4' must be changed from 3 to 0,...
242
+ c = c.set_categories([4, 3, 2, 1])
243
+
244
+ # positions are changed
245
+ tm.assert_numpy_array_equal(c._codes, np.array([3, 2, 1, 0, 3], dtype=np.int8))
246
+
247
+ # categories are now in new order
248
+ tm.assert_index_equal(c.categories, Index([4, 3, 2, 1]))
249
+
250
+ # output is the same
251
+ exp = np.array([1, 2, 3, 4, 1], dtype=np.int64)
252
+ tm.assert_numpy_array_equal(np.asarray(c), exp)
253
+ assert c.min() == 4
254
+ assert c.max() == 1
255
+
256
+ # set_categories should set the ordering if specified
257
+ c2 = c.set_categories([4, 3, 2, 1], ordered=False)
258
+ assert not c2.ordered
259
+
260
+ tm.assert_numpy_array_equal(np.asarray(c), np.asarray(c2))
261
+
262
+ # set_categories should pass thru the ordering
263
+ c2 = c.set_ordered(False).set_categories([4, 3, 2, 1])
264
+ assert not c2.ordered
265
+
266
+ tm.assert_numpy_array_equal(np.asarray(c), np.asarray(c2))
267
+
268
+ @pytest.mark.parametrize(
269
+ "values, categories, new_categories",
270
+ [
271
+ # No NaNs, same cats, same order
272
+ (["a", "b", "a"], ["a", "b"], ["a", "b"]),
273
+ # No NaNs, same cats, different order
274
+ (["a", "b", "a"], ["a", "b"], ["b", "a"]),
275
+ # Same, unsorted
276
+ (["b", "a", "a"], ["a", "b"], ["a", "b"]),
277
+ # No NaNs, same cats, different order
278
+ (["b", "a", "a"], ["a", "b"], ["b", "a"]),
279
+ # NaNs
280
+ (["a", "b", "c"], ["a", "b"], ["a", "b"]),
281
+ (["a", "b", "c"], ["a", "b"], ["b", "a"]),
282
+ (["b", "a", "c"], ["a", "b"], ["a", "b"]),
283
+ (["b", "a", "c"], ["a", "b"], ["b", "a"]),
284
+ # Introduce NaNs
285
+ (["a", "b", "c"], ["a", "b"], ["a"]),
286
+ (["a", "b", "c"], ["a", "b"], ["b"]),
287
+ (["b", "a", "c"], ["a", "b"], ["a"]),
288
+ (["b", "a", "c"], ["a", "b"], ["b"]),
289
+ # No overlap
290
+ (["a", "b", "c"], ["a", "b"], ["d", "e"]),
291
+ ],
292
+ )
293
+ def test_set_categories_many(self, values, categories, new_categories, ordered):
294
+ msg = "Constructing a Categorical with a dtype and values containing"
295
+
296
+ warn1 = Pandas4Warning if set(values).difference(categories) else None
297
+ with tm.assert_produces_warning(warn1, match=msg):
298
+ c = Categorical(values, categories)
299
+
300
+ warn2 = Pandas4Warning if set(values).difference(new_categories) else None
301
+ with tm.assert_produces_warning(warn2, match=msg):
302
+ expected = Categorical(values, new_categories, ordered)
303
+
304
+ result = c.set_categories(new_categories, ordered=ordered)
305
+ tm.assert_categorical_equal(result, expected)
306
+
307
+ def test_set_categories_rename_less(self):
308
+ # GH 24675
309
+ cat = Categorical(["A", "B"])
310
+ result = cat.set_categories(["A"], rename=True)
311
+ expected = Categorical(["A", np.nan])
312
+ tm.assert_categorical_equal(result, expected)
313
+
314
+ def test_set_categories_private(self):
315
+ cat = Categorical(["a", "b", "c"], categories=["a", "b", "c", "d"])
316
+ cat._set_categories(["a", "c", "d", "e"])
317
+ expected = Categorical(["a", "c", "d"], categories=list("acde"))
318
+ tm.assert_categorical_equal(cat, expected)
319
+
320
+ # fastpath
321
+ cat = Categorical(["a", "b", "c"], categories=["a", "b", "c", "d"])
322
+ cat._set_categories(["a", "c", "d", "e"], fastpath=True)
323
+ expected = Categorical(["a", "c", "d"], categories=list("acde"))
324
+ tm.assert_categorical_equal(cat, expected)
325
+
326
+ def test_remove_categories(self):
327
+ cat = Categorical(["a", "b", "c", "a"], ordered=True)
328
+ old = cat.copy()
329
+ new = Categorical(["a", "b", np.nan, "a"], categories=["a", "b"], ordered=True)
330
+
331
+ res = cat.remove_categories("c")
332
+ tm.assert_categorical_equal(cat, old)
333
+ tm.assert_categorical_equal(res, new)
334
+
335
+ res = cat.remove_categories(["c"])
336
+ tm.assert_categorical_equal(cat, old)
337
+ tm.assert_categorical_equal(res, new)
338
+
339
+ @pytest.mark.parametrize("removals", [["c"], ["c", np.nan], "c", ["c", "c"]])
340
+ def test_remove_categories_raises(self, removals):
341
+ cat = Categorical(["a", "b", "a"])
342
+ message = re.escape("removals must all be in old categories: {'c'}")
343
+
344
+ with pytest.raises(ValueError, match=message):
345
+ cat.remove_categories(removals)
346
+
347
+ def test_remove_unused_categories(self):
348
+ c = Categorical(["a", "b", "c", "d", "a"], categories=["a", "b", "c", "d", "e"])
349
+ exp_categories_all = Index(["a", "b", "c", "d", "e"])
350
+ exp_categories_dropped = Index(["a", "b", "c", "d"])
351
+
352
+ tm.assert_index_equal(c.categories, exp_categories_all)
353
+
354
+ res = c.remove_unused_categories()
355
+ tm.assert_index_equal(res.categories, exp_categories_dropped)
356
+ tm.assert_index_equal(c.categories, exp_categories_all)
357
+
358
+ # with NaN values (GH11599)
359
+ c = Categorical(["a", "b", "c", np.nan], categories=["a", "b", "c", "d", "e"])
360
+ res = c.remove_unused_categories()
361
+ tm.assert_index_equal(res.categories, Index(np.array(["a", "b", "c"])))
362
+ exp_codes = np.array([0, 1, 2, -1], dtype=np.int8)
363
+ tm.assert_numpy_array_equal(res.codes, exp_codes)
364
+ tm.assert_index_equal(c.categories, exp_categories_all)
365
+
366
+ val = ["F", np.nan, "D", "B", "D", "F", np.nan]
367
+ cat = Categorical(values=val, categories=list("ABCDEFG"))
368
+ out = cat.remove_unused_categories()
369
+ tm.assert_index_equal(out.categories, Index(["B", "D", "F"]))
370
+ exp_codes = np.array([2, -1, 1, 0, 1, 2, -1], dtype=np.int8)
371
+ tm.assert_numpy_array_equal(out.codes, exp_codes)
372
+ assert out.tolist() == val
373
+
374
+ alpha = list("abcdefghijklmnopqrstuvwxyz")
375
+ val = np.random.default_rng(2).choice(alpha[::2], 10000).astype("object")
376
+ val[np.random.default_rng(2).choice(len(val), 100)] = np.nan
377
+
378
+ cat = Categorical(values=val, categories=alpha)
379
+ out = cat.remove_unused_categories()
380
+ assert out.tolist() == val.tolist()
381
+
382
+
383
+ class TestCategoricalAPIWithFactor:
384
+ def test_describe(self):
385
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
386
+ # string type
387
+ desc = factor.describe()
388
+ assert factor.ordered
389
+ exp_index = CategoricalIndex(
390
+ ["a", "b", "c"], name="categories", ordered=factor.ordered
391
+ )
392
+ expected = DataFrame(
393
+ {"counts": [3, 2, 3], "freqs": [3 / 8.0, 2 / 8.0, 3 / 8.0]}, index=exp_index
394
+ )
395
+ tm.assert_frame_equal(desc, expected)
396
+
397
+ # check unused categories
398
+ cat = factor.copy()
399
+ cat = cat.set_categories(["a", "b", "c", "d"])
400
+ desc = cat.describe()
401
+
402
+ exp_index = CategoricalIndex(
403
+ list("abcd"), ordered=factor.ordered, name="categories"
404
+ )
405
+ expected = DataFrame(
406
+ {"counts": [3, 2, 3, 0], "freqs": [3 / 8.0, 2 / 8.0, 3 / 8.0, 0]},
407
+ index=exp_index,
408
+ )
409
+ tm.assert_frame_equal(desc, expected)
410
+
411
+ # check an integer one
412
+ cat = Categorical([1, 2, 3, 1, 2, 3, 3, 2, 1, 1, 1])
413
+ desc = cat.describe()
414
+ exp_index = CategoricalIndex([1, 2, 3], ordered=cat.ordered, name="categories")
415
+ expected = DataFrame(
416
+ {"counts": [5, 3, 3], "freqs": [5 / 11.0, 3 / 11.0, 3 / 11.0]},
417
+ index=exp_index,
418
+ )
419
+ tm.assert_frame_equal(desc, expected)
420
+
421
+ # https://github.com/pandas-dev/pandas/issues/3678
422
+ # describe should work with NaN
423
+ cat = Categorical([np.nan, 1, 2, 2])
424
+ desc = cat.describe()
425
+ expected = DataFrame(
426
+ {"counts": [1, 2, 1], "freqs": [1 / 4.0, 2 / 4.0, 1 / 4.0]},
427
+ index=CategoricalIndex(
428
+ [1, 2, np.nan], categories=[1, 2], name="categories"
429
+ ),
430
+ )
431
+ tm.assert_frame_equal(desc, expected)
432
+
433
+
434
+ class TestPrivateCategoricalAPI:
435
+ def test_codes_immutable(self):
436
+ # Codes should be read only
437
+ c = Categorical(["a", "b", "c", "a", np.nan])
438
+ exp = np.array([0, 1, 2, 0, -1], dtype="int8")
439
+ tm.assert_numpy_array_equal(c.codes, exp)
440
+
441
+ # Assignments to codes should raise
442
+ msg = "property 'codes' of 'Categorical' object has no setter"
443
+ with pytest.raises(AttributeError, match=msg):
444
+ c.codes = np.array([0, 1, 2, 0, 1], dtype="int8")
445
+
446
+ # changes in the codes array should raise
447
+ codes = c.codes
448
+
449
+ with pytest.raises(ValueError, match="assignment destination is read-only"):
450
+ codes[4] = 1
451
+
452
+ # But even after getting the codes, the original array should still be
453
+ # writeable!
454
+ c[4] = "a"
455
+ exp = np.array([0, 1, 2, 0, 0], dtype="int8")
456
+ tm.assert_numpy_array_equal(c.codes, exp)
457
+ c._codes[4] = 2
458
+ exp = np.array([0, 1, 2, 0, 2], dtype="int8")
459
+ tm.assert_numpy_array_equal(c.codes, exp)
460
+
461
+ @pytest.mark.parametrize(
462
+ "codes, old, new, expected",
463
+ [
464
+ ([0, 1], ["a", "b"], ["a", "b"], [0, 1]),
465
+ ([0, 1], ["b", "a"], ["b", "a"], [0, 1]),
466
+ ([0, 1], ["a", "b"], ["b", "a"], [1, 0]),
467
+ ([0, 1], ["b", "a"], ["a", "b"], [1, 0]),
468
+ ([0, 1, 0, 1], ["a", "b"], ["a", "b", "c"], [0, 1, 0, 1]),
469
+ ([0, 1, 2, 2], ["a", "b", "c"], ["a", "b"], [0, 1, -1, -1]),
470
+ ([0, 1, -1], ["a", "b", "c"], ["a", "b", "c"], [0, 1, -1]),
471
+ ([0, 1, -1], ["a", "b", "c"], ["b"], [-1, 0, -1]),
472
+ ([0, 1, -1], ["a", "b", "c"], ["d"], [-1, -1, -1]),
473
+ ([0, 1, -1], ["a", "b", "c"], [], [-1, -1, -1]),
474
+ ([-1, -1], [], ["a", "b"], [-1, -1]),
475
+ ([1, 0], ["b", "a"], ["a", "b"], [0, 1]),
476
+ ],
477
+ )
478
+ def test_recode_to_categories(self, codes, old, new, expected):
479
+ codes = np.asanyarray(codes, dtype=np.int8)
480
+ expected = np.asanyarray(expected, dtype=np.int8)
481
+ old = Index(old)
482
+ new = Index(new)
483
+ result = recode_for_categories(codes, old, new, copy=True)
484
+ tm.assert_numpy_array_equal(result, expected)
485
+
486
+ def test_recode_to_categories_large(self):
487
+ N = 1000
488
+ codes = np.arange(N)
489
+ old = Index(codes)
490
+ expected = np.arange(N - 1, -1, -1, dtype=np.int16)
491
+ new = Index(expected)
492
+ result = recode_for_categories(codes, old, new, copy=True)
493
+ tm.assert_numpy_array_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_astype.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas.errors import Pandas4Warning
5
+
6
+ from pandas import (
7
+ NA,
8
+ Categorical,
9
+ CategoricalDtype,
10
+ CategoricalIndex,
11
+ DatetimeIndex,
12
+ Interval,
13
+ NaT,
14
+ Period,
15
+ Timestamp,
16
+ array,
17
+ to_datetime,
18
+ )
19
+ import pandas._testing as tm
20
+
21
+
22
+ class TestAstype:
23
+ @pytest.mark.parametrize("cls", [Categorical, CategoricalIndex])
24
+ @pytest.mark.parametrize("values", [[1, np.nan], [Timestamp("2000"), NaT]])
25
+ def test_astype_nan_to_int(self, cls, values):
26
+ # GH#28406
27
+ obj = cls(values)
28
+
29
+ msg = "Cannot (cast|convert)"
30
+ with pytest.raises((ValueError, TypeError), match=msg):
31
+ obj.astype(int)
32
+
33
+ @pytest.mark.parametrize(
34
+ "expected",
35
+ [
36
+ array(["2019", "2020"], dtype="datetime64[ns, UTC]"),
37
+ array([0, 0], dtype="timedelta64[ns]"),
38
+ array([Period("2019"), Period("2020")], dtype="period[Y-DEC]"),
39
+ array([Interval(0, 1), Interval(1, 2)], dtype="interval"),
40
+ array([1, NA], dtype="Int64"),
41
+ ],
42
+ )
43
+ def test_astype_category_to_extension_dtype(self, expected):
44
+ # GH#28668
45
+ result = expected.astype("category").astype(expected.dtype)
46
+
47
+ tm.assert_extension_array_equal(result, expected)
48
+
49
+ @pytest.mark.parametrize(
50
+ "dtype, expected",
51
+ [
52
+ (
53
+ "datetime64[ns]",
54
+ np.array(["2015-01-01T00:00:00.000000000"], dtype="datetime64[ns]"),
55
+ ),
56
+ (
57
+ "datetime64[ns, MET]",
58
+ DatetimeIndex([Timestamp("2015-01-01 00:00:00+0100", tz="MET")]).array,
59
+ ),
60
+ ],
61
+ )
62
+ def test_astype_to_datetime64(self, dtype, expected):
63
+ # GH#28448
64
+ result = Categorical(["2015-01-01"]).astype(dtype)
65
+ assert result == expected
66
+
67
+ def test_astype_str_int_categories_to_nullable_int(self):
68
+ # GH#39616
69
+ dtype = CategoricalDtype([str(i) for i in range(5)])
70
+ codes = np.random.default_rng(2).integers(5, size=20)
71
+ arr = Categorical.from_codes(codes, dtype=dtype)
72
+
73
+ res = arr.astype("Int64")
74
+ expected = array(codes, dtype="Int64")
75
+ tm.assert_extension_array_equal(res, expected)
76
+
77
+ def test_astype_str_int_categories_to_nullable_float(self):
78
+ # GH#39616
79
+ dtype = CategoricalDtype([str(i / 2) for i in range(5)])
80
+ codes = np.random.default_rng(2).integers(5, size=20)
81
+ arr = Categorical.from_codes(codes, dtype=dtype)
82
+
83
+ res = arr.astype("Float64")
84
+ expected = array(codes, dtype="Float64") / 2
85
+ tm.assert_extension_array_equal(res, expected)
86
+
87
+ def test_astype(self, ordered):
88
+ # string
89
+ cat = Categorical(list("abbaaccc"), ordered=ordered)
90
+ result = cat.astype(object)
91
+ expected = np.array(cat)
92
+ tm.assert_numpy_array_equal(result, expected)
93
+
94
+ msg = r"Cannot cast object|str dtype to float64"
95
+ with pytest.raises(ValueError, match=msg):
96
+ cat.astype(float)
97
+
98
+ # numeric
99
+ cat = Categorical([0, 1, 2, 2, 1, 0, 1, 0, 2], ordered=ordered)
100
+ result = cat.astype(object)
101
+ expected = np.array(cat, dtype=object)
102
+ tm.assert_numpy_array_equal(result, expected)
103
+
104
+ result = cat.astype(int)
105
+ expected = np.array(cat, dtype="int")
106
+ tm.assert_numpy_array_equal(result, expected)
107
+
108
+ result = cat.astype(float)
109
+ expected = np.array(cat, dtype=float)
110
+ tm.assert_numpy_array_equal(result, expected)
111
+
112
+ @pytest.mark.parametrize("dtype_ordered", [True, False])
113
+ def test_astype_category(self, dtype_ordered, ordered):
114
+ # GH#10696/GH#18593
115
+ data = list("abcaacbab")
116
+ cat = Categorical(data, categories=list("bac"), ordered=ordered)
117
+
118
+ # standard categories
119
+ dtype = CategoricalDtype(ordered=dtype_ordered)
120
+ result = cat.astype(dtype)
121
+ expected = Categorical(data, categories=cat.categories, ordered=dtype_ordered)
122
+ tm.assert_categorical_equal(result, expected)
123
+
124
+ # non-standard categories
125
+ dtype = CategoricalDtype(list("adc"), dtype_ordered)
126
+ msg = "Constructing a Categorical with a dtype and values containing"
127
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
128
+ result = cat.astype(dtype)
129
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
130
+ expected = Categorical(data, dtype=dtype)
131
+ tm.assert_categorical_equal(result, expected)
132
+
133
+ if dtype_ordered is False:
134
+ # dtype='category' can't specify ordered, so only test once
135
+ result = cat.astype("category")
136
+ expected = cat
137
+ tm.assert_categorical_equal(result, expected)
138
+
139
+ def test_astype_category_copy_false_nocopy_codes(self):
140
+ # GH#62000
141
+ cat = Categorical([3, 2, 4, 1])
142
+ new = cat.astype("category", copy=False)
143
+ assert tm.shares_memory(new.codes, cat.codes)
144
+ new = cat.astype("category", copy=True)
145
+ assert not tm.shares_memory(new.codes, cat.codes)
146
+
147
+ def test_astype_object_datetime_categories(self):
148
+ # GH#40754
149
+ cat = Categorical(to_datetime(["2021-03-27", NaT]))
150
+ result = cat.astype(object)
151
+ expected = np.array([Timestamp("2021-03-27 00:00:00"), NaT], dtype="object")
152
+ tm.assert_numpy_array_equal(result, expected)
153
+
154
+ def test_astype_object_timestamp_categories(self):
155
+ # GH#18024
156
+ cat = Categorical([Timestamp("2014-01-01")])
157
+ result = cat.astype(object)
158
+ expected = np.array([Timestamp("2014-01-01 00:00:00")], dtype="object")
159
+ tm.assert_numpy_array_equal(result, expected)
160
+
161
+ def test_astype_category_readonly_mask_values(self):
162
+ # GH#53658
163
+ arr = array([0, 1, 2], dtype="Int64")
164
+ arr._mask.flags["WRITEABLE"] = False
165
+ result = arr.astype("category")
166
+ expected = array([0, 1, 2], dtype="Int64").astype("category")
167
+ tm.assert_extension_array_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_constructors.py ADDED
@@ -0,0 +1,836 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datetime import (
2
+ date,
3
+ datetime,
4
+ )
5
+
6
+ import numpy as np
7
+ import pytest
8
+
9
+ from pandas.errors import Pandas4Warning
10
+
11
+ from pandas.core.dtypes.common import (
12
+ is_float_dtype,
13
+ is_integer_dtype,
14
+ )
15
+ from pandas.core.dtypes.dtypes import CategoricalDtype
16
+
17
+ import pandas as pd
18
+ from pandas import (
19
+ Categorical,
20
+ CategoricalIndex,
21
+ DatetimeIndex,
22
+ Index,
23
+ Interval,
24
+ IntervalIndex,
25
+ MultiIndex,
26
+ NaT,
27
+ RangeIndex,
28
+ Series,
29
+ Timestamp,
30
+ date_range,
31
+ period_range,
32
+ timedelta_range,
33
+ )
34
+ import pandas._testing as tm
35
+
36
+
37
+ class TestCategoricalConstructors:
38
+ def test_categorical_from_cat_and_dtype_str_preserve_ordered(self):
39
+ # GH#49309 we should preserve orderedness in `res`
40
+ cat = Categorical([3, 1], categories=[3, 2, 1], ordered=True)
41
+
42
+ res = Categorical(cat, dtype="category")
43
+ assert res.dtype.ordered
44
+
45
+ def test_categorical_disallows_scalar(self):
46
+ # GH#38433
47
+ with pytest.raises(TypeError, match="Categorical input must be list-like"):
48
+ Categorical("A", categories=["A", "B"])
49
+
50
+ def test_categorical_1d_only(self):
51
+ # ndim > 1
52
+ msg = "> 1 ndim Categorical are not supported at this time"
53
+ with pytest.raises(NotImplementedError, match=msg):
54
+ Categorical(np.array([list("abcd")]))
55
+
56
+ def test_validate_ordered(self):
57
+ # see gh-14058
58
+ exp_msg = "'ordered' must either be 'True' or 'False'"
59
+ exp_err = TypeError
60
+
61
+ # This should be a boolean.
62
+ ordered = np.array([0, 1, 2])
63
+
64
+ with pytest.raises(exp_err, match=exp_msg):
65
+ Categorical([1, 2, 3], ordered=ordered)
66
+
67
+ with pytest.raises(exp_err, match=exp_msg):
68
+ Categorical.from_codes(
69
+ [0, 0, 1], categories=["a", "b", "c"], ordered=ordered
70
+ )
71
+
72
+ def test_constructor_empty(self):
73
+ # GH 17248
74
+ c = Categorical([])
75
+ expected = Index([])
76
+ tm.assert_index_equal(c.categories, expected)
77
+
78
+ c = Categorical([], categories=[1, 2, 3])
79
+ expected = Index([1, 2, 3], dtype=np.int64)
80
+ tm.assert_index_equal(c.categories, expected)
81
+
82
+ def test_constructor_empty_boolean(self):
83
+ # see gh-22702
84
+ cat = Categorical([], categories=[True, False])
85
+ categories = sorted(cat.categories.tolist())
86
+ assert categories == [False, True]
87
+
88
+ def test_constructor_tuples(self):
89
+ values = np.array([(1,), (1, 2), (1,), (1, 2)], dtype=object)
90
+ result = Categorical(values)
91
+ expected = Index([(1,), (1, 2)], tupleize_cols=False)
92
+ tm.assert_index_equal(result.categories, expected)
93
+ assert result.ordered is False
94
+
95
+ def test_constructor_tuples_datetimes(self):
96
+ # numpy will auto reshape when all of the tuples are the
97
+ # same len, so add an extra one with 2 items and slice it off
98
+ values = np.array(
99
+ [
100
+ (Timestamp("2010-01-01"),),
101
+ (Timestamp("2010-01-02"),),
102
+ (Timestamp("2010-01-01"),),
103
+ (Timestamp("2010-01-02"),),
104
+ ("a", "b"),
105
+ ],
106
+ dtype=object,
107
+ )[:-1]
108
+ result = Categorical(values)
109
+ expected = Index(
110
+ [(Timestamp("2010-01-01"),), (Timestamp("2010-01-02"),)],
111
+ tupleize_cols=False,
112
+ )
113
+ tm.assert_index_equal(result.categories, expected)
114
+
115
+ def test_constructor_unsortable(self):
116
+ # it works!
117
+ arr = np.array([1, 2, 3, datetime.now()], dtype="O")
118
+ factor = Categorical(arr, ordered=False)
119
+ assert not factor.ordered
120
+
121
+ # this however will raise as cannot be sorted
122
+ msg = (
123
+ "'values' is not ordered, please explicitly specify the "
124
+ "categories order by passing in a categories argument."
125
+ )
126
+ with pytest.raises(TypeError, match=msg):
127
+ Categorical(arr, ordered=True)
128
+
129
+ def test_constructor_interval(self):
130
+ result = Categorical(
131
+ [Interval(1, 2), Interval(2, 3), Interval(3, 6)], ordered=True
132
+ )
133
+ ii = IntervalIndex([Interval(1, 2), Interval(2, 3), Interval(3, 6)])
134
+ exp = Categorical(ii, ordered=True)
135
+ tm.assert_categorical_equal(result, exp)
136
+ tm.assert_index_equal(result.categories, ii)
137
+
138
+ def test_constructor(self):
139
+ exp_arr = np.array(["a", "b", "c", "a", "b", "c"], dtype=np.object_)
140
+ c1 = Categorical(exp_arr)
141
+ tm.assert_numpy_array_equal(c1.__array__(), exp_arr)
142
+ c2 = Categorical(exp_arr, categories=["a", "b", "c"])
143
+ tm.assert_numpy_array_equal(c2.__array__(), exp_arr)
144
+ c2 = Categorical(exp_arr, categories=["c", "b", "a"])
145
+ tm.assert_numpy_array_equal(c2.__array__(), exp_arr)
146
+
147
+ # categories must be unique
148
+ msg = "Categorical categories must be unique"
149
+ with pytest.raises(ValueError, match=msg):
150
+ Categorical([1, 2], [1, 2, 2])
151
+
152
+ with pytest.raises(ValueError, match=msg):
153
+ Categorical(["a", "b"], ["a", "b", "b"])
154
+
155
+ # The default should be unordered
156
+ c1 = Categorical(["a", "b", "c", "a"])
157
+ assert not c1.ordered
158
+
159
+ # Categorical as input
160
+ c1 = Categorical(["a", "b", "c", "a"])
161
+ c2 = Categorical(c1)
162
+ tm.assert_categorical_equal(c1, c2)
163
+
164
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "b", "c", "d"])
165
+ c2 = Categorical(c1)
166
+ tm.assert_categorical_equal(c1, c2)
167
+
168
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "c", "b"])
169
+ c2 = Categorical(c1)
170
+ tm.assert_categorical_equal(c1, c2)
171
+
172
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "c", "b"])
173
+ c2 = Categorical(c1, categories=["a", "b", "c"])
174
+ tm.assert_numpy_array_equal(c1.__array__(), c2.__array__())
175
+ tm.assert_index_equal(c2.categories, Index(["a", "b", "c"]))
176
+
177
+ # Series of dtype category
178
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "b", "c", "d"])
179
+ c2 = Categorical(Series(c1))
180
+ tm.assert_categorical_equal(c1, c2)
181
+
182
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "c", "b"])
183
+ c2 = Categorical(Series(c1))
184
+ tm.assert_categorical_equal(c1, c2)
185
+
186
+ # Series
187
+ c1 = Categorical(["a", "b", "c", "a"])
188
+ c2 = Categorical(Series(["a", "b", "c", "a"]))
189
+ tm.assert_categorical_equal(c1, c2)
190
+
191
+ c1 = Categorical(["a", "b", "c", "a"], categories=["a", "b", "c", "d"])
192
+ c2 = Categorical(Series(["a", "b", "c", "a"]), categories=["a", "b", "c", "d"])
193
+ tm.assert_categorical_equal(c1, c2)
194
+
195
+ # This should result in integer categories, not float!
196
+ cat = Categorical([1, 2, 3, np.nan], categories=[1, 2, 3])
197
+ assert is_integer_dtype(cat.categories)
198
+
199
+ # https://github.com/pandas-dev/pandas/issues/3678
200
+ cat = Categorical([np.nan, 1, 2, 3])
201
+ assert is_integer_dtype(cat.categories)
202
+
203
+ # this should result in floats
204
+ cat = Categorical([np.nan, 1, 2.0, 3])
205
+ assert is_float_dtype(cat.categories)
206
+
207
+ cat = Categorical([np.nan, 1.0, 2.0, 3.0])
208
+ assert is_float_dtype(cat.categories)
209
+
210
+ # This doesn't work -> this would probably need some kind of "remember
211
+ # the original type" feature to try to cast the array interface result
212
+ # to...
213
+
214
+ # vals = np.asarray(cat[cat.notna()])
215
+ # assert is_integer_dtype(vals)
216
+
217
+ # corner cases
218
+ cat = Categorical([1])
219
+ assert len(cat.categories) == 1
220
+ assert cat.categories[0] == 1
221
+ assert len(cat.codes) == 1
222
+ assert cat.codes[0] == 0
223
+
224
+ cat = Categorical(["a"])
225
+ assert len(cat.categories) == 1
226
+ assert cat.categories[0] == "a"
227
+ assert len(cat.codes) == 1
228
+ assert cat.codes[0] == 0
229
+
230
+ # two arrays
231
+ # - when the first is an integer dtype and the second is not
232
+ # - when the resulting codes are all -1/NaN
233
+ msg = "Constructing a Categorical with a dtype and values containing"
234
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
235
+ Categorical([0, 1, 2, 0, 1, 2], categories=["a", "b", "c"])
236
+
237
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
238
+ Categorical([0, 1, 2, 0, 1, 2], categories=[3, 4, 5])
239
+
240
+ # the next one are from the old docs
241
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
242
+ Categorical([0, 1, 2, 0, 1, 2], [1, 2, 3])
243
+ cat = Categorical([1, 2], categories=[1, 2, 3])
244
+
245
+ # this is a legitimate constructor
246
+ with tm.assert_produces_warning(None):
247
+ Categorical(np.array([], dtype="int64"), categories=[3, 2, 1], ordered=True)
248
+
249
+ def test_constructor_with_existing_categories(self):
250
+ # GH25318: constructing with pd.Series used to bogusly skip recoding
251
+ # categories
252
+ c0 = Categorical(["a", "b", "c", "a"])
253
+ msg = "Constructing a Categorical with a dtype and values containing"
254
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
255
+ c1 = Categorical(["a", "b", "c", "a"], categories=["b", "c"])
256
+
257
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
258
+ c2 = Categorical(c0, categories=c1.categories)
259
+ tm.assert_categorical_equal(c1, c2)
260
+
261
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
262
+ c3 = Categorical(Series(c0), categories=c1.categories)
263
+ tm.assert_categorical_equal(c1, c3)
264
+
265
+ def test_constructor_not_sequence(self):
266
+ # https://github.com/pandas-dev/pandas/issues/16022
267
+ msg = r"^Parameter 'categories' must be list-like, was"
268
+ with pytest.raises(TypeError, match=msg):
269
+ Categorical(["a", "b"], categories="a")
270
+
271
+ def test_constructor_with_null(self):
272
+ # Cannot have NaN in categories
273
+ msg = "Categorical categories cannot be null"
274
+ with pytest.raises(ValueError, match=msg):
275
+ Categorical([np.nan, "a", "b", "c"], categories=[np.nan, "a", "b", "c"])
276
+
277
+ with pytest.raises(ValueError, match=msg):
278
+ Categorical([None, "a", "b", "c"], categories=[None, "a", "b", "c"])
279
+
280
+ with pytest.raises(ValueError, match=msg):
281
+ Categorical(
282
+ DatetimeIndex(["nat", "20160101"]),
283
+ categories=[NaT, Timestamp("20160101")],
284
+ )
285
+
286
+ def test_constructor_with_index(self):
287
+ ci = CategoricalIndex(list("aabbca"), categories=list("cab"))
288
+ tm.assert_categorical_equal(ci.values, Categorical(ci))
289
+
290
+ ci = CategoricalIndex(list("aabbca"), categories=list("cab"))
291
+ tm.assert_categorical_equal(
292
+ ci.values, Categorical(ci.astype(object), categories=ci.categories)
293
+ )
294
+
295
+ def test_constructor_with_generator(self):
296
+ # This was raising an Error in isna(single_val).any() because isna
297
+ # returned a scalar for a generator
298
+
299
+ exp = Categorical([0, 1, 2])
300
+ cat = Categorical(x for x in [0, 1, 2])
301
+ tm.assert_categorical_equal(cat, exp)
302
+ cat = Categorical(range(3))
303
+ tm.assert_categorical_equal(cat, exp)
304
+
305
+ MultiIndex.from_product([range(5), ["a", "b", "c"]])
306
+
307
+ # check that categories accept generators and sequences
308
+ cat = Categorical([0, 1, 2], categories=(x for x in [0, 1, 2]))
309
+ tm.assert_categorical_equal(cat, exp)
310
+ cat = Categorical([0, 1, 2], categories=range(3))
311
+ tm.assert_categorical_equal(cat, exp)
312
+
313
+ def test_constructor_with_rangeindex(self):
314
+ # RangeIndex is preserved in Categories
315
+ rng = Index(range(3))
316
+
317
+ cat = Categorical(rng)
318
+ tm.assert_index_equal(cat.categories, rng, exact=True)
319
+
320
+ cat = Categorical([1, 2, 0], categories=rng)
321
+ tm.assert_index_equal(cat.categories, rng, exact=True)
322
+
323
+ @pytest.mark.parametrize(
324
+ "dtl",
325
+ [
326
+ date_range("1995-01-01 00:00:00", periods=5, freq="s"),
327
+ date_range("1995-01-01 00:00:00", periods=5, freq="s", tz="US/Eastern"),
328
+ timedelta_range("1 day", periods=5, freq="s"),
329
+ ],
330
+ )
331
+ def test_constructor_with_datetimelike(self, dtl):
332
+ # see gh-12077
333
+ # constructor with a datetimelike and NaT
334
+
335
+ s = Series(dtl)
336
+ c = Categorical(s)
337
+
338
+ expected = type(dtl)(s)
339
+ expected._data.freq = None
340
+
341
+ tm.assert_index_equal(c.categories, expected)
342
+ tm.assert_numpy_array_equal(c.codes, np.arange(5, dtype="int8"))
343
+
344
+ # with NaT
345
+ s2 = s.copy()
346
+ s2.iloc[-1] = NaT
347
+ c = Categorical(s2)
348
+
349
+ expected = type(dtl)(s2.dropna())
350
+ expected._data.freq = None
351
+
352
+ tm.assert_index_equal(c.categories, expected)
353
+
354
+ exp = np.array([0, 1, 2, 3, -1], dtype=np.int8)
355
+ tm.assert_numpy_array_equal(c.codes, exp)
356
+
357
+ result = repr(c)
358
+ assert "NaT" in result
359
+
360
+ def test_constructor_from_index_series_datetimetz(self):
361
+ idx = date_range("2015-01-01 10:00", freq="D", periods=3, tz="US/Eastern")
362
+ idx = idx._with_freq(None) # freq not preserved in result.categories
363
+ result = Categorical(idx)
364
+ tm.assert_index_equal(result.categories, idx)
365
+
366
+ result = Categorical(Series(idx))
367
+ tm.assert_index_equal(result.categories, idx)
368
+
369
+ def test_constructor_date_objects(self):
370
+ # we dont cast date objects to timestamps, matching Index constructor
371
+ v = date.today()
372
+
373
+ cat = Categorical([v, v])
374
+ assert cat.categories.dtype == object
375
+ assert type(cat.categories[0]) is date
376
+
377
+ def test_constructor_from_index_series_timedelta(self):
378
+ idx = timedelta_range("1 days", freq="D", periods=3)
379
+ idx = idx._with_freq(None) # freq not preserved in result.categories
380
+ result = Categorical(idx)
381
+ tm.assert_index_equal(result.categories, idx)
382
+
383
+ result = Categorical(Series(idx))
384
+ tm.assert_index_equal(result.categories, idx)
385
+
386
+ def test_constructor_from_index_series_period(self):
387
+ idx = period_range("2015-01-01", freq="D", periods=3)
388
+ result = Categorical(idx)
389
+ tm.assert_index_equal(result.categories, idx)
390
+
391
+ result = Categorical(Series(idx))
392
+ tm.assert_index_equal(result.categories, idx)
393
+
394
+ @pytest.mark.parametrize(
395
+ "values",
396
+ [
397
+ np.array([1.0, 1.2, 1.8, np.nan]),
398
+ np.array([1, 2, 3], dtype="int64"),
399
+ ["a", "b", "c", np.nan],
400
+ [pd.Period("2014-01"), pd.Period("2014-02"), NaT],
401
+ [Timestamp("2014-01-01"), Timestamp("2014-01-02"), NaT],
402
+ [
403
+ Timestamp("2014-01-01", tz="US/Eastern"),
404
+ Timestamp("2014-01-02", tz="US/Eastern"),
405
+ NaT,
406
+ ],
407
+ ],
408
+ )
409
+ def test_constructor_invariant(self, values):
410
+ # GH 14190
411
+ c = Categorical(values)
412
+ c2 = Categorical(c)
413
+ tm.assert_categorical_equal(c, c2)
414
+
415
+ @pytest.mark.parametrize("ordered", [True, False])
416
+ def test_constructor_with_dtype(self, ordered):
417
+ categories = ["b", "a", "c"]
418
+ dtype = CategoricalDtype(categories, ordered=ordered)
419
+ result = Categorical(["a", "b", "a", "c"], dtype=dtype)
420
+ expected = Categorical(
421
+ ["a", "b", "a", "c"], categories=categories, ordered=ordered
422
+ )
423
+ tm.assert_categorical_equal(result, expected)
424
+ assert result.ordered is ordered
425
+
426
+ def test_constructor_dtype_and_others_raises(self):
427
+ dtype = CategoricalDtype(["a", "b"], ordered=True)
428
+ msg = "Cannot specify `categories` or `ordered` together with `dtype`."
429
+ with pytest.raises(ValueError, match=msg):
430
+ Categorical(["a", "b"], categories=["a", "b"], dtype=dtype)
431
+
432
+ with pytest.raises(ValueError, match=msg):
433
+ Categorical(["a", "b"], ordered=True, dtype=dtype)
434
+
435
+ with pytest.raises(ValueError, match=msg):
436
+ Categorical(["a", "b"], ordered=False, dtype=dtype)
437
+
438
+ @pytest.mark.parametrize("categories", [None, ["a", "b"], ["a", "c"]])
439
+ def test_constructor_str_category(self, categories, ordered):
440
+ warn = Pandas4Warning if categories == ["a", "c"] else None
441
+ msg = "Constructing a Categorical with a dtype and values containing"
442
+ with tm.assert_produces_warning(warn, match=msg):
443
+ result = Categorical(
444
+ ["a", "b"], categories=categories, ordered=ordered, dtype="category"
445
+ )
446
+ expected = Categorical(["a", "b"], categories=categories, ordered=ordered)
447
+ tm.assert_categorical_equal(result, expected)
448
+
449
+ def test_constructor_str_unknown(self):
450
+ with pytest.raises(ValueError, match="Unknown dtype"):
451
+ Categorical([1, 2], dtype="foo")
452
+
453
+ def test_constructor_np_strs(self):
454
+ # GH#31499 Hashtable.map_locations needs to work on np.str_ objects
455
+ # We can't pass all-strings because the constructor would cast
456
+ # those to StringDtype post-PDEP14
457
+ cat = Categorical(["1", "0", "1", 2], [np.str_("0"), np.str_("1"), 2])
458
+ assert all(isinstance(x, (np.str_, int)) for x in cat.categories)
459
+
460
+ def test_constructor_from_categorical_with_dtype(self):
461
+ dtype = CategoricalDtype(["a", "b", "c"], ordered=True)
462
+ values = Categorical(["a", "b", "d"])
463
+ msg = "Constructing a Categorical with a dtype and values containing"
464
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
465
+ result = Categorical(values, dtype=dtype)
466
+ # We use dtype.categories, not values.categories
467
+ expected = Categorical(
468
+ ["a", "b", None], categories=["a", "b", "c"], ordered=True
469
+ )
470
+ tm.assert_categorical_equal(result, expected)
471
+
472
+ def test_constructor_from_categorical_with_unknown_dtype(self):
473
+ dtype = CategoricalDtype(None, ordered=True)
474
+ values = Categorical(["a", "b", "d"])
475
+ result = Categorical(values, dtype=dtype)
476
+ # We use values.categories, not dtype.categories
477
+ expected = Categorical(
478
+ ["a", "b", "d"], categories=["a", "b", "d"], ordered=True
479
+ )
480
+ tm.assert_categorical_equal(result, expected)
481
+
482
+ def test_constructor_from_categorical_string(self):
483
+ values = Categorical(["a", "b", "d"])
484
+ # use categories, ordered
485
+ msg = "Constructing a Categorical with a dtype and values containing"
486
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
487
+ result = Categorical(
488
+ values, categories=["a", "b", "c"], ordered=True, dtype="category"
489
+ )
490
+ expected = Categorical(
491
+ ["a", "b", None], categories=["a", "b", "c"], ordered=True
492
+ )
493
+ tm.assert_categorical_equal(result, expected)
494
+
495
+ # No string
496
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
497
+ result = Categorical(values, categories=["a", "b", "c"], ordered=True)
498
+ tm.assert_categorical_equal(result, expected)
499
+
500
+ def test_constructor_with_categorical_categories(self):
501
+ # GH17884
502
+ expected = Categorical(["a", "b"], categories=["a", "b", "c"])
503
+
504
+ result = Categorical(["a", "b"], categories=Categorical(["a", "b", "c"]))
505
+ tm.assert_categorical_equal(result, expected)
506
+
507
+ result = Categorical(["a", "b"], categories=CategoricalIndex(["a", "b", "c"]))
508
+ tm.assert_categorical_equal(result, expected)
509
+
510
+ @pytest.mark.parametrize("klass", [lambda x: np.array(x, dtype=object), list])
511
+ def test_construction_with_null(self, klass, nulls_fixture):
512
+ # https://github.com/pandas-dev/pandas/issues/31927
513
+ values = klass(["a", nulls_fixture, "b"])
514
+ result = Categorical(values)
515
+
516
+ dtype = CategoricalDtype(["a", "b"])
517
+ codes = [0, -1, 1]
518
+ expected = Categorical.from_codes(codes=codes, dtype=dtype)
519
+
520
+ tm.assert_categorical_equal(result, expected)
521
+
522
+ @pytest.mark.parametrize("validate", [True, False])
523
+ def test_from_codes_nullable_int_categories(self, any_numeric_ea_dtype, validate):
524
+ # GH#39649
525
+ cats = pd.array(range(5), dtype=any_numeric_ea_dtype)
526
+ codes = np.random.default_rng(2).integers(5, size=3)
527
+ dtype = CategoricalDtype(cats)
528
+ arr = Categorical.from_codes(codes, dtype=dtype, validate=validate)
529
+ assert arr.categories.dtype == cats.dtype
530
+ tm.assert_index_equal(arr.categories, Index(cats))
531
+
532
+ def test_from_codes_empty(self):
533
+ cat = ["a", "b", "c"]
534
+ result = Categorical.from_codes([], categories=cat)
535
+ expected = Categorical([], categories=cat)
536
+
537
+ tm.assert_categorical_equal(result, expected)
538
+
539
+ @pytest.mark.parametrize("validate", [True, False])
540
+ def test_from_codes_validate(self, validate):
541
+ # GH53122
542
+ dtype = CategoricalDtype(["a", "b"])
543
+ if validate:
544
+ with pytest.raises(ValueError, match="codes need to be between "):
545
+ Categorical.from_codes([4, 5], dtype=dtype, validate=validate)
546
+ else:
547
+ # passes, though has incorrect codes, but that's the user responsibility
548
+ Categorical.from_codes([4, 5], dtype=dtype, validate=validate)
549
+
550
+ def test_from_codes_too_few_categories(self):
551
+ dtype = CategoricalDtype(categories=[1, 2])
552
+ msg = "codes need to be between "
553
+ with pytest.raises(ValueError, match=msg):
554
+ Categorical.from_codes([1, 2], categories=dtype.categories)
555
+ with pytest.raises(ValueError, match=msg):
556
+ Categorical.from_codes([1, 2], dtype=dtype)
557
+
558
+ def test_from_codes_non_int_codes(self):
559
+ dtype = CategoricalDtype(categories=[1, 2])
560
+ msg = "codes need to be array-like integers"
561
+ with pytest.raises(ValueError, match=msg):
562
+ Categorical.from_codes(["a"], categories=dtype.categories)
563
+ with pytest.raises(ValueError, match=msg):
564
+ Categorical.from_codes(["a"], dtype=dtype)
565
+
566
+ def test_from_codes_non_unique_categories(self):
567
+ with pytest.raises(ValueError, match="Categorical categories must be unique"):
568
+ Categorical.from_codes([0, 1, 2], categories=["a", "a", "b"])
569
+
570
+ def test_from_codes_nan_cat_included(self):
571
+ with pytest.raises(ValueError, match="Categorical categories cannot be null"):
572
+ Categorical.from_codes([0, 1, 2], categories=["a", "b", np.nan])
573
+
574
+ def test_from_codes_too_negative(self):
575
+ dtype = CategoricalDtype(categories=["a", "b", "c"])
576
+ msg = r"codes need to be between -1 and len\(categories\)-1"
577
+ with pytest.raises(ValueError, match=msg):
578
+ Categorical.from_codes([-2, 1, 2], categories=dtype.categories)
579
+ with pytest.raises(ValueError, match=msg):
580
+ Categorical.from_codes([-2, 1, 2], dtype=dtype)
581
+
582
+ def test_from_codes(self):
583
+ dtype = CategoricalDtype(categories=["a", "b", "c"])
584
+ exp = Categorical(["a", "b", "c"], ordered=False)
585
+ res = Categorical.from_codes([0, 1, 2], categories=dtype.categories)
586
+ tm.assert_categorical_equal(exp, res)
587
+
588
+ res = Categorical.from_codes([0, 1, 2], dtype=dtype)
589
+ tm.assert_categorical_equal(exp, res)
590
+
591
+ @pytest.mark.parametrize("klass", [Categorical, CategoricalIndex])
592
+ def test_from_codes_with_categorical_categories(self, klass):
593
+ # GH17884
594
+ expected = Categorical(["a", "b"], categories=["a", "b", "c"])
595
+
596
+ result = Categorical.from_codes([0, 1], categories=klass(["a", "b", "c"]))
597
+ tm.assert_categorical_equal(result, expected)
598
+
599
+ @pytest.mark.parametrize("klass", [Categorical, CategoricalIndex])
600
+ def test_from_codes_with_non_unique_categorical_categories(self, klass):
601
+ with pytest.raises(ValueError, match="Categorical categories must be unique"):
602
+ Categorical.from_codes([0, 1], klass(["a", "b", "a"]))
603
+
604
+ def test_from_codes_with_nan_code(self):
605
+ # GH21767
606
+ codes = [1, 2, np.nan]
607
+ dtype = CategoricalDtype(categories=["a", "b", "c"])
608
+ with pytest.raises(ValueError, match="codes need to be array-like integers"):
609
+ Categorical.from_codes(codes, categories=dtype.categories)
610
+ with pytest.raises(ValueError, match="codes need to be array-like integers"):
611
+ Categorical.from_codes(codes, dtype=dtype)
612
+
613
+ @pytest.mark.parametrize("codes", [[1.0, 2.0, 0], [1.1, 2.0, 0]])
614
+ def test_from_codes_with_float(self, codes):
615
+ # GH21767
616
+ # float codes should raise even if values are equal to integers
617
+ dtype = CategoricalDtype(categories=["a", "b", "c"])
618
+
619
+ msg = "codes need to be array-like integers"
620
+ with pytest.raises(ValueError, match=msg):
621
+ Categorical.from_codes(codes, dtype.categories)
622
+ with pytest.raises(ValueError, match=msg):
623
+ Categorical.from_codes(codes, dtype=dtype)
624
+
625
+ def test_from_codes_with_dtype_raises(self):
626
+ msg = "Cannot specify"
627
+ with pytest.raises(ValueError, match=msg):
628
+ Categorical.from_codes(
629
+ [0, 1], categories=["a", "b"], dtype=CategoricalDtype(["a", "b"])
630
+ )
631
+
632
+ with pytest.raises(ValueError, match=msg):
633
+ Categorical.from_codes(
634
+ [0, 1], ordered=True, dtype=CategoricalDtype(["a", "b"])
635
+ )
636
+
637
+ def test_from_codes_neither(self):
638
+ msg = "Both were None"
639
+ with pytest.raises(ValueError, match=msg):
640
+ Categorical.from_codes([0, 1])
641
+
642
+ def test_from_codes_with_nullable_int(self):
643
+ codes = pd.array([0, 1], dtype="Int64")
644
+ categories = ["a", "b"]
645
+
646
+ result = Categorical.from_codes(codes, categories=categories)
647
+ expected = Categorical.from_codes(codes.to_numpy(int), categories=categories)
648
+
649
+ tm.assert_categorical_equal(result, expected)
650
+
651
+ def test_from_codes_with_nullable_int_na_raises(self):
652
+ codes = pd.array([0, None], dtype="Int64")
653
+ categories = ["a", "b"]
654
+
655
+ msg = "codes cannot contain NA values"
656
+ with pytest.raises(ValueError, match=msg):
657
+ Categorical.from_codes(codes, categories=categories)
658
+
659
+ @pytest.mark.parametrize("dtype", [None, "category"])
660
+ def test_from_inferred_categories(self, dtype):
661
+ cats = ["a", "b"]
662
+ codes = np.array([0, 0, 1, 1], dtype="i8")
663
+ result = Categorical._from_inferred_categories(cats, codes, dtype)
664
+ expected = Categorical.from_codes(codes, cats)
665
+ tm.assert_categorical_equal(result, expected)
666
+
667
+ @pytest.mark.parametrize("dtype", [None, "category"])
668
+ def test_from_inferred_categories_sorts(self, dtype):
669
+ cats = ["b", "a"]
670
+ codes = np.array([0, 1, 1, 1], dtype="i8")
671
+ result = Categorical._from_inferred_categories(cats, codes, dtype)
672
+ expected = Categorical.from_codes([1, 0, 0, 0], ["a", "b"])
673
+ tm.assert_categorical_equal(result, expected)
674
+
675
+ def test_from_inferred_categories_dtype(self):
676
+ cats = ["a", "b", "d"]
677
+ codes = np.array([0, 1, 0, 2], dtype="i8")
678
+ dtype = CategoricalDtype(["c", "b", "a"], ordered=True)
679
+ msg = "Constructing a Categorical with a dtype and values containing"
680
+ with tm.assert_produces_warning(
681
+ Pandas4Warning, match=msg, check_stacklevel=False
682
+ ):
683
+ result = Categorical._from_inferred_categories(cats, codes, dtype)
684
+ expected = Categorical(
685
+ ["a", "b", "a", None], categories=["c", "b", "a"], ordered=True
686
+ )
687
+ tm.assert_categorical_equal(result, expected)
688
+
689
+ def test_from_inferred_categories_coerces(self):
690
+ cats = ["1", "2", "bad"]
691
+ codes = np.array([0, 0, 1, 2], dtype="i8")
692
+ dtype = CategoricalDtype([1, 2])
693
+ msg = "Constructing a Categorical with a dtype and values containing"
694
+ with tm.assert_produces_warning(
695
+ Pandas4Warning, match=msg, check_stacklevel=False
696
+ ):
697
+ result = Categorical._from_inferred_categories(cats, codes, dtype)
698
+ expected = Categorical([1, 1, 2, np.nan])
699
+ tm.assert_categorical_equal(result, expected)
700
+
701
+ def test_construction_with_ordered(self, ordered):
702
+ # GH 9347, 9190
703
+ cat = Categorical([0, 1, 2], ordered=ordered)
704
+ assert cat.ordered == bool(ordered)
705
+
706
+ def test_constructor_imaginary(self):
707
+ values = [1, 2, 3 + 1j]
708
+ c1 = Categorical(values)
709
+ tm.assert_index_equal(c1.categories, Index(values))
710
+ tm.assert_numpy_array_equal(np.array(c1), np.array(values))
711
+
712
+ def test_constructor_string_and_tuples(self):
713
+ # GH 21416
714
+ c = Categorical(np.array(["c", ("a", "b"), ("b", "a"), "c"], dtype=object))
715
+ expected_index = Index([("a", "b"), ("b", "a"), "c"])
716
+ assert c.categories.equals(expected_index)
717
+
718
+ def test_interval(self):
719
+ idx = pd.interval_range(0, 10, periods=10)
720
+ cat = Categorical(idx, categories=idx)
721
+ expected_codes = np.arange(10, dtype="int8")
722
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
723
+ tm.assert_index_equal(cat.categories, idx)
724
+
725
+ # infer categories
726
+ cat = Categorical(idx)
727
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
728
+ tm.assert_index_equal(cat.categories, idx)
729
+
730
+ # list values
731
+ cat = Categorical(list(idx))
732
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
733
+ tm.assert_index_equal(cat.categories, idx)
734
+
735
+ # list values, categories
736
+ cat = Categorical(list(idx), categories=list(idx))
737
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
738
+ tm.assert_index_equal(cat.categories, idx)
739
+
740
+ # shuffled
741
+ values = idx.take([1, 2, 0])
742
+ cat = Categorical(values, categories=idx)
743
+ tm.assert_numpy_array_equal(cat.codes, np.array([1, 2, 0], dtype="int8"))
744
+ tm.assert_index_equal(cat.categories, idx)
745
+
746
+ # extra
747
+ values = pd.interval_range(8, 11, periods=3)
748
+ msg = "Constructing a Categorical with a dtype and values containing"
749
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
750
+ cat = Categorical(values, categories=idx)
751
+ expected_codes = np.array([8, 9, -1], dtype="int8")
752
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
753
+ tm.assert_index_equal(cat.categories, idx)
754
+
755
+ # overlapping
756
+ idx = IntervalIndex([Interval(0, 2), Interval(0, 1)])
757
+ cat = Categorical(idx, categories=idx)
758
+ expected_codes = np.array([0, 1], dtype="int8")
759
+ tm.assert_numpy_array_equal(cat.codes, expected_codes)
760
+ tm.assert_index_equal(cat.categories, idx)
761
+
762
+ def test_categorical_extension_array_nullable(self, nulls_fixture):
763
+ # GH:
764
+ arr = pd.array([nulls_fixture] * 2, dtype=pd.StringDtype())
765
+ result = Categorical(arr)
766
+ assert arr.dtype == result.categories.dtype
767
+ expected = Categorical(Series([pd.NA, pd.NA], dtype=arr.dtype))
768
+ tm.assert_categorical_equal(result, expected)
769
+
770
+ def test_from_sequence_copy(self):
771
+ cat = Categorical(np.arange(5).repeat(2))
772
+ result = Categorical._from_sequence(cat, dtype=cat.dtype, copy=False)
773
+
774
+ # more generally, we'd be OK with a view
775
+ assert result._codes is cat._codes
776
+
777
+ result = Categorical._from_sequence(cat, dtype=cat.dtype, copy=True)
778
+
779
+ assert not tm.shares_memory(result, cat)
780
+
781
+ def test_constructor_datetime64_non_nano(self):
782
+ categories = np.arange(10).view("M8[D]")
783
+ values = categories[::2].copy()
784
+
785
+ cat = Categorical(values, categories=categories)
786
+ assert (cat == values).all()
787
+
788
+ def test_constructor_preserves_freq(self):
789
+ # GH33830 freq retention in categorical
790
+ dti = date_range("2016-01-01", periods=5)
791
+
792
+ expected = dti.freq
793
+
794
+ cat = Categorical(dti)
795
+ result = cat.categories.freq
796
+
797
+ assert expected == result
798
+
799
+ @pytest.mark.parametrize(
800
+ "values, categories",
801
+ [
802
+ [range(5), None],
803
+ [range(4), range(5)],
804
+ [[0, 1, 2, 3], range(5)],
805
+ [[], range(5)],
806
+ ],
807
+ )
808
+ def test_range_values_preserves_rangeindex_categories(self, values, categories):
809
+ result = Categorical(values=values, categories=categories).categories
810
+ expected = RangeIndex(range(5))
811
+ tm.assert_index_equal(result, expected, exact=True)
812
+
813
+ def test_categorical_preserve_object_dtype_from_pandas(self, using_infer_string):
814
+ # GH#61778
815
+ pylist = ["foo", "bar", "baz"]
816
+ ser = Series(pylist, dtype="object")
817
+ idx = Index(pylist, dtype="object")
818
+ arr = np.array(pylist, dtype="object")
819
+
820
+ cat_from_ser = Categorical(ser)
821
+ cat_from_idx = Categorical(idx)
822
+ cat_from_arr = Categorical(arr)
823
+ cat_from_list = Categorical(pylist)
824
+
825
+ # Series/Index with object dtype: infer string
826
+ # dtype if all elements are strings
827
+ assert cat_from_ser.categories.dtype == object
828
+ assert cat_from_idx.categories.dtype == object
829
+
830
+ if using_infer_string:
831
+ # Numpy array or list: infer string dtype
832
+ assert cat_from_arr.categories.dtype == "str"
833
+ assert cat_from_list.categories.dtype == "str"
834
+ else:
835
+ assert cat_from_arr.categories.dtype == object
836
+ assert cat_from_list.categories.dtype == object
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_dtypes.py ADDED
@@ -0,0 +1,147 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas.errors import Pandas4Warning
5
+
6
+ from pandas.core.dtypes.dtypes import CategoricalDtype
7
+
8
+ from pandas import (
9
+ Categorical,
10
+ CategoricalIndex,
11
+ Index,
12
+ IntervalIndex,
13
+ Series,
14
+ Timestamp,
15
+ )
16
+ import pandas._testing as tm
17
+
18
+
19
+ class TestCategoricalDtypes:
20
+ def test_categories_match_up_to_permutation(self):
21
+ # test dtype comparisons between cats
22
+
23
+ c1 = Categorical(list("aabca"), categories=list("abc"), ordered=False)
24
+ c2 = Categorical(list("aabca"), categories=list("cab"), ordered=False)
25
+ c3 = Categorical(list("aabca"), categories=list("cab"), ordered=True)
26
+ assert c1._categories_match_up_to_permutation(c1)
27
+ assert c2._categories_match_up_to_permutation(c2)
28
+ assert c3._categories_match_up_to_permutation(c3)
29
+ assert c1._categories_match_up_to_permutation(c2)
30
+ assert not c1._categories_match_up_to_permutation(c3)
31
+ assert not c1._categories_match_up_to_permutation(Index(list("aabca")))
32
+ assert not c1._categories_match_up_to_permutation(c1.astype(object))
33
+ assert c1._categories_match_up_to_permutation(CategoricalIndex(c1))
34
+ assert c1._categories_match_up_to_permutation(
35
+ CategoricalIndex(c1, categories=list("cab"))
36
+ )
37
+ assert not c1._categories_match_up_to_permutation(
38
+ CategoricalIndex(c1, ordered=True)
39
+ )
40
+
41
+ # GH 16659
42
+ s1 = Series(c1)
43
+ s2 = Series(c2)
44
+ s3 = Series(c3)
45
+ assert c1._categories_match_up_to_permutation(s1)
46
+ assert c2._categories_match_up_to_permutation(s2)
47
+ assert c3._categories_match_up_to_permutation(s3)
48
+ assert c1._categories_match_up_to_permutation(s2)
49
+ assert not c1._categories_match_up_to_permutation(s3)
50
+ assert not c1._categories_match_up_to_permutation(s1.astype(object))
51
+
52
+ def test_set_dtype_same(self):
53
+ c = Categorical(["a", "b", "c"])
54
+ result = c._set_dtype(CategoricalDtype(["a", "b", "c"]), copy=True)
55
+ tm.assert_categorical_equal(result, c)
56
+
57
+ def test_set_dtype_new_categories(self):
58
+ c = Categorical(["a", "b", "c"])
59
+ result = c._set_dtype(CategoricalDtype(list("abcd")), copy=True)
60
+ tm.assert_numpy_array_equal(result.codes, c.codes)
61
+ tm.assert_index_equal(result.dtype.categories, Index(list("abcd")))
62
+
63
+ @pytest.mark.parametrize(
64
+ "values, categories, new_categories, warn",
65
+ [
66
+ # No NaNs, same cats, same order
67
+ (["a", "b", "a"], ["a", "b"], ["a", "b"], None),
68
+ # No NaNs, same cats, different order
69
+ (["a", "b", "a"], ["a", "b"], ["b", "a"], None),
70
+ # Same, unsorted
71
+ (["b", "a", "a"], ["a", "b"], ["a", "b"], None),
72
+ # No NaNs, same cats, different order
73
+ (["b", "a", "a"], ["a", "b"], ["b", "a"], None),
74
+ # NaNs
75
+ (["a", "b", "c"], ["a", "b"], ["a", "b"], None),
76
+ (["a", "b", "c"], ["a", "b"], ["b", "a"], None),
77
+ (["b", "a", "c"], ["a", "b"], ["a", "b"], None),
78
+ (["b", "a", "c"], ["a", "b"], ["b", "a"], None),
79
+ # Introduce NaNs
80
+ (["a", "b", "c"], ["a", "b"], ["a"], Pandas4Warning),
81
+ (["a", "b", "c"], ["a", "b"], ["b"], Pandas4Warning),
82
+ (["b", "a", "c"], ["a", "b"], ["a"], Pandas4Warning),
83
+ (["b", "a", "c"], ["a", "b"], ["b"], Pandas4Warning),
84
+ # No overlap
85
+ (["a", "b", "c"], ["a", "b"], ["d", "e"], Pandas4Warning),
86
+ ],
87
+ )
88
+ def test_set_dtype_many(self, values, categories, new_categories, warn, ordered):
89
+ msg = "Constructing a Categorical with a dtype and values containing"
90
+ warn1 = Pandas4Warning if set(values).difference(categories) else None
91
+ with tm.assert_produces_warning(warn1, match=msg):
92
+ c = Categorical(values, categories)
93
+ warn2 = Pandas4Warning if set(values).difference(new_categories) else None
94
+ with tm.assert_produces_warning(warn2, match=msg):
95
+ expected = Categorical(values, new_categories, ordered)
96
+ result = c._set_dtype(expected.dtype, copy=True)
97
+ tm.assert_categorical_equal(result, expected)
98
+
99
+ def test_set_dtype_no_overlap(self):
100
+ msg = "Constructing a Categorical with a dtype and values containing"
101
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
102
+ c = Categorical(["a", "b", "c"], ["d", "e"])
103
+ result = c._set_dtype(CategoricalDtype(["a", "b"]), copy=True)
104
+ expected = Categorical([None, None, None], categories=["a", "b"])
105
+ tm.assert_categorical_equal(result, expected)
106
+
107
+ def test_codes_dtypes(self):
108
+ # GH 8453
109
+ result = Categorical(["foo", "bar", "baz"])
110
+ assert result.codes.dtype == "int8"
111
+
112
+ result = Categorical([f"foo{i:05d}" for i in range(400)])
113
+ assert result.codes.dtype == "int16"
114
+
115
+ result = Categorical([f"foo{i:05d}" for i in range(40000)])
116
+ assert result.codes.dtype == "int32"
117
+
118
+ # adding cats
119
+ result = Categorical(["foo", "bar", "baz"])
120
+ assert result.codes.dtype == "int8"
121
+ result = result.add_categories([f"foo{i:05d}" for i in range(400)])
122
+ assert result.codes.dtype == "int16"
123
+
124
+ # removing cats
125
+ result = result.remove_categories([f"foo{i:05d}" for i in range(300)])
126
+ assert result.codes.dtype == "int8"
127
+
128
+ def test_iter_python_types(self):
129
+ # GH-19909
130
+ cat = Categorical([1, 2])
131
+ assert isinstance(next(iter(cat)), int)
132
+ assert isinstance(cat.tolist()[0], int)
133
+
134
+ def test_iter_python_types_datetime(self):
135
+ cat = Categorical([Timestamp("2017-01-01"), Timestamp("2017-01-02")])
136
+ assert isinstance(next(iter(cat)), Timestamp)
137
+ assert isinstance(cat.tolist()[0], Timestamp)
138
+
139
+ def test_interval_index_category(self):
140
+ # GH 38316
141
+ index = IntervalIndex.from_breaks(np.arange(3, dtype="uint64"))
142
+
143
+ result = CategoricalIndex(index).dtype.categories
144
+ expected = IntervalIndex.from_arrays(
145
+ [0, 1], [1, 2], dtype="interval[uint64, right]"
146
+ )
147
+ tm.assert_index_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_indexing.py ADDED
@@ -0,0 +1,386 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ from pandas import (
7
+ NA,
8
+ Categorical,
9
+ CategoricalIndex,
10
+ Index,
11
+ Interval,
12
+ IntervalIndex,
13
+ NaT,
14
+ PeriodIndex,
15
+ Series,
16
+ Timedelta,
17
+ Timestamp,
18
+ )
19
+ import pandas._testing as tm
20
+ import pandas.core.common as com
21
+
22
+
23
+ class TestCategoricalIndexingWithFactor:
24
+ def test_getitem(self):
25
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
26
+ assert factor[0] == "a"
27
+ assert factor[-1] == "c"
28
+
29
+ subf = factor[[0, 1, 2]]
30
+ tm.assert_numpy_array_equal(subf._codes, np.array([0, 1, 1], dtype=np.int8))
31
+
32
+ subf = factor[np.asarray(factor) == "c"]
33
+ tm.assert_numpy_array_equal(subf._codes, np.array([2, 2, 2], dtype=np.int8))
34
+
35
+ def test_setitem(self):
36
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
37
+ # int/positional
38
+ c = factor.copy()
39
+ c[0] = "b"
40
+ assert c[0] == "b"
41
+ c[-1] = "a"
42
+ assert c[-1] == "a"
43
+
44
+ # boolean
45
+ c = factor.copy()
46
+ indexer = np.zeros(len(c), dtype="bool")
47
+ indexer[0] = True
48
+ indexer[-1] = True
49
+ c[indexer] = "c"
50
+ expected = Categorical(["c", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
51
+
52
+ tm.assert_categorical_equal(c, expected)
53
+
54
+ @pytest.mark.parametrize("categories", [None, ["b", "a"]])
55
+ def test_setitem_same_but_unordered(self, categories):
56
+ # GH-24142
57
+ other = Categorical(["b", "a"], categories=categories)
58
+ target = Categorical(["a", "b"], categories=["a", "b"])
59
+ mask = np.array([True, False])
60
+ target[mask] = other[mask]
61
+ expected = Categorical(["b", "b"], categories=["a", "b"])
62
+ tm.assert_categorical_equal(target, expected)
63
+
64
+ @pytest.mark.parametrize(
65
+ "other",
66
+ [
67
+ Categorical(["b", "a"], categories=["b", "a", "c"]),
68
+ Categorical(["b", "a"], categories=["a", "b", "c"]),
69
+ Categorical(["a", "a"], categories=["a"]),
70
+ Categorical(["b", "b"], categories=["b"]),
71
+ ],
72
+ )
73
+ def test_setitem_different_unordered_raises(self, other):
74
+ # GH-24142
75
+ target = Categorical(["a", "b"], categories=["a", "b"])
76
+ mask = np.array([True, False])
77
+ msg = "Cannot set a Categorical with another, without identical categories"
78
+ with pytest.raises(TypeError, match=msg):
79
+ target[mask] = other[mask]
80
+
81
+ @pytest.mark.parametrize(
82
+ "other",
83
+ [
84
+ Categorical(["b", "a"]),
85
+ Categorical(["b", "a"], categories=["b", "a"], ordered=True),
86
+ Categorical(["b", "a"], categories=["a", "b", "c"], ordered=True),
87
+ ],
88
+ )
89
+ def test_setitem_same_ordered_raises(self, other):
90
+ # Gh-24142
91
+ target = Categorical(["a", "b"], categories=["a", "b"], ordered=True)
92
+ mask = np.array([True, False])
93
+ msg = "Cannot set a Categorical with another, without identical categories"
94
+ with pytest.raises(TypeError, match=msg):
95
+ target[mask] = other[mask]
96
+
97
+ def test_setitem_tuple(self):
98
+ # GH#20439
99
+ cat = Categorical([(0, 1), (0, 2), (0, 1)])
100
+
101
+ # This should not raise
102
+ cat[1] = cat[0]
103
+ assert cat[1] == (0, 1)
104
+
105
+ def test_setitem_listlike(self):
106
+ # GH#9469
107
+ # properly coerce the input indexers
108
+
109
+ cat = Categorical(
110
+ np.random.default_rng(2).integers(0, 5, size=150000).astype(np.int8)
111
+ ).add_categories([-1000])
112
+ indexer = np.array([100000]).astype(np.int64)
113
+ cat[indexer] = -1000
114
+
115
+ # we are asserting the code result here
116
+ # which maps to the -1000 category
117
+ result = cat.codes[np.array([100000]).astype(np.int64)]
118
+ tm.assert_numpy_array_equal(result, np.array([5], dtype="int8"))
119
+
120
+
121
+ class TestCategoricalIndexing:
122
+ def test_getitem_slice(self):
123
+ cat = Categorical(["a", "b", "c", "d", "a", "b", "c"])
124
+ sliced = cat[3]
125
+ assert sliced == "d"
126
+
127
+ sliced = cat[3:5]
128
+ expected = Categorical(["d", "a"], categories=["a", "b", "c", "d"])
129
+ tm.assert_categorical_equal(sliced, expected)
130
+
131
+ def test_getitem_listlike(self):
132
+ # GH 9469
133
+ # properly coerce the input indexers
134
+
135
+ c = Categorical(
136
+ np.random.default_rng(2).integers(0, 5, size=150000).astype(np.int8)
137
+ )
138
+ result = c.codes[np.array([100000]).astype(np.int64)]
139
+ expected = c[np.array([100000]).astype(np.int64)].codes
140
+ tm.assert_numpy_array_equal(result, expected)
141
+
142
+ def test_periodindex(self):
143
+ idx1 = PeriodIndex(
144
+ ["2014-01", "2014-01", "2014-02", "2014-02", "2014-03", "2014-03"],
145
+ freq="M",
146
+ )
147
+
148
+ cat1 = Categorical(idx1)
149
+ str(cat1)
150
+ exp_arr = np.array([0, 0, 1, 1, 2, 2], dtype=np.int8)
151
+ exp_idx = PeriodIndex(["2014-01", "2014-02", "2014-03"], freq="M")
152
+ tm.assert_numpy_array_equal(cat1._codes, exp_arr)
153
+ tm.assert_index_equal(cat1.categories, exp_idx)
154
+
155
+ idx2 = PeriodIndex(
156
+ ["2014-03", "2014-03", "2014-02", "2014-01", "2014-03", "2014-01"],
157
+ freq="M",
158
+ )
159
+ cat2 = Categorical(idx2, ordered=True)
160
+ str(cat2)
161
+ exp_arr = np.array([2, 2, 1, 0, 2, 0], dtype=np.int8)
162
+ exp_idx2 = PeriodIndex(["2014-01", "2014-02", "2014-03"], freq="M")
163
+ tm.assert_numpy_array_equal(cat2._codes, exp_arr)
164
+ tm.assert_index_equal(cat2.categories, exp_idx2)
165
+
166
+ idx3 = PeriodIndex(
167
+ [
168
+ "2013-12",
169
+ "2013-11",
170
+ "2013-10",
171
+ "2013-09",
172
+ "2013-08",
173
+ "2013-07",
174
+ "2013-05",
175
+ ],
176
+ freq="M",
177
+ )
178
+ cat3 = Categorical(idx3, ordered=True)
179
+ exp_arr = np.array([6, 5, 4, 3, 2, 1, 0], dtype=np.int8)
180
+ exp_idx = PeriodIndex(
181
+ [
182
+ "2013-05",
183
+ "2013-07",
184
+ "2013-08",
185
+ "2013-09",
186
+ "2013-10",
187
+ "2013-11",
188
+ "2013-12",
189
+ ],
190
+ freq="M",
191
+ )
192
+ tm.assert_numpy_array_equal(cat3._codes, exp_arr)
193
+ tm.assert_index_equal(cat3.categories, exp_idx)
194
+
195
+ @pytest.mark.parametrize(
196
+ "null_val",
197
+ [None, np.nan, NaT, NA, math.nan, "NaT", "nat", "NAT", "nan", "NaN", "NAN"],
198
+ )
199
+ def test_periodindex_on_null_types(self, null_val):
200
+ # GH 46673
201
+ result = PeriodIndex(["2022-04-06", "2022-04-07", null_val], freq="D")
202
+ expected = PeriodIndex(["2022-04-06", "2022-04-07", "NaT"], dtype="period[D]")
203
+ assert result[2] is NaT
204
+ tm.assert_index_equal(result, expected)
205
+
206
+ @pytest.mark.parametrize("new_categories", [[1, 2, 3, 4], [1, 2]])
207
+ def test_categories_assignments_wrong_length_raises(self, new_categories):
208
+ cat = Categorical(["a", "b", "c", "a"])
209
+ msg = (
210
+ "new categories need to have the same number of items "
211
+ "as the old categories!"
212
+ )
213
+ with pytest.raises(ValueError, match=msg):
214
+ cat.rename_categories(new_categories)
215
+
216
+ # Combinations of sorted/unique:
217
+ @pytest.mark.parametrize(
218
+ "idx_values", [[1, 2, 3, 4], [1, 3, 2, 4], [1, 3, 3, 4], [1, 2, 2, 4]]
219
+ )
220
+ # Combinations of missing/unique
221
+ @pytest.mark.parametrize("key_values", [[1, 2], [1, 5], [1, 1], [5, 5]])
222
+ @pytest.mark.parametrize("key_class", [Categorical, CategoricalIndex])
223
+ @pytest.mark.parametrize("dtype", [None, "category", "key"])
224
+ def test_get_indexer_non_unique(self, idx_values, key_values, key_class, dtype):
225
+ # GH 21448
226
+ key = key_class(key_values, categories=range(1, 6))
227
+
228
+ if dtype == "key":
229
+ dtype = key.dtype
230
+
231
+ # Test for flat index and CategoricalIndex with same/different cats:
232
+ idx = Index(idx_values, dtype=dtype)
233
+ expected, exp_miss = idx.get_indexer_non_unique(key_values)
234
+ result, res_miss = idx.get_indexer_non_unique(key)
235
+
236
+ tm.assert_numpy_array_equal(expected, result)
237
+ tm.assert_numpy_array_equal(exp_miss, res_miss)
238
+
239
+ exp_unique = idx.unique().get_indexer(key_values)
240
+ res_unique = idx.unique().get_indexer(key)
241
+ tm.assert_numpy_array_equal(res_unique, exp_unique)
242
+
243
+ def test_where_unobserved_nan(self):
244
+ ser = Series(Categorical(["a", "b"]))
245
+ result = ser.where([True, False])
246
+ expected = Series(Categorical(["a", None], categories=["a", "b"]))
247
+ tm.assert_series_equal(result, expected)
248
+
249
+ # all NA
250
+ ser = Series(Categorical(["a", "b"]))
251
+ result = ser.where([False, False])
252
+ expected = Series(Categorical([None, None], categories=["a", "b"]))
253
+ tm.assert_series_equal(result, expected)
254
+
255
+ def test_where_unobserved_categories(self):
256
+ ser = Series(Categorical(["a", "b", "c"], categories=["d", "c", "b", "a"]))
257
+ result = ser.where([True, True, False], other="b")
258
+ expected = Series(Categorical(["a", "b", "b"], categories=ser.cat.categories))
259
+ tm.assert_series_equal(result, expected)
260
+
261
+ def test_where_other_categorical(self):
262
+ ser = Series(Categorical(["a", "b", "c"], categories=["d", "c", "b", "a"]))
263
+ other = Categorical(["b", "c", "a"], categories=["a", "c", "b", "d"])
264
+ result = ser.where([True, False, True], other)
265
+ expected = Series(Categorical(["a", "c", "c"], dtype=ser.dtype))
266
+ tm.assert_series_equal(result, expected)
267
+
268
+ def test_where_new_category_raises(self):
269
+ ser = Series(Categorical(["a", "b", "c"]))
270
+ msg = "Cannot setitem on a Categorical with a new category"
271
+ with pytest.raises(TypeError, match=msg):
272
+ ser.where([True, False, True], "d")
273
+
274
+ def test_where_ordered_differs_rasies(self):
275
+ ser = Series(
276
+ Categorical(["a", "b", "c"], categories=["d", "c", "b", "a"], ordered=True)
277
+ )
278
+ other = Categorical(
279
+ ["b", "c", "a"], categories=["a", "c", "b", "d"], ordered=True
280
+ )
281
+ with pytest.raises(TypeError, match="without identical categories"):
282
+ ser.where([True, False, True], other)
283
+
284
+
285
+ class TestContains:
286
+ def test_contains(self):
287
+ # GH#21508
288
+ cat = Categorical(list("aabbca"), categories=list("cab"))
289
+
290
+ assert "b" in cat
291
+ assert "z" not in cat
292
+ assert np.nan not in cat
293
+ with pytest.raises(TypeError, match="unhashable type: 'list'"):
294
+ assert [1] in cat
295
+
296
+ # assert codes NOT in index
297
+ assert 0 not in cat
298
+ assert 1 not in cat
299
+
300
+ cat = Categorical([*list("aabbca"), np.nan], categories=list("cab"))
301
+ assert np.nan in cat
302
+
303
+ @pytest.mark.parametrize(
304
+ "item, expected",
305
+ [
306
+ (Interval(0, 1), True),
307
+ (1.5, True),
308
+ (Interval(0.5, 1.5), False),
309
+ ("a", False),
310
+ (Timestamp(1), False),
311
+ (Timedelta(1), False),
312
+ ],
313
+ ids=str,
314
+ )
315
+ def test_contains_interval(self, item, expected):
316
+ # GH#23705
317
+ cat = Categorical(IntervalIndex.from_breaks(range(3)))
318
+ result = item in cat
319
+ assert result is expected
320
+
321
+ def test_contains_list(self):
322
+ # GH#21729
323
+ cat = Categorical([1, 2, 3])
324
+
325
+ assert "a" not in cat
326
+
327
+ with pytest.raises(TypeError, match="unhashable type"):
328
+ ["a"] in cat
329
+
330
+ with pytest.raises(TypeError, match="unhashable type"):
331
+ ["a", "b"] in cat
332
+
333
+
334
+ @pytest.mark.parametrize("index", [True, False])
335
+ def test_mask_with_boolean(index):
336
+ ser = Series(range(3))
337
+ idx = Categorical([True, False, True])
338
+ if index:
339
+ idx = CategoricalIndex(idx)
340
+
341
+ assert com.is_bool_indexer(idx)
342
+ result = ser[idx]
343
+ expected = ser[idx.astype("object")]
344
+ tm.assert_series_equal(result, expected)
345
+
346
+
347
+ @pytest.mark.parametrize("index", [True, False])
348
+ def test_mask_with_boolean_na_treated_as_false(index):
349
+ # https://github.com/pandas-dev/pandas/issues/31503
350
+ ser = Series(range(3))
351
+ idx = Categorical([True, False, None])
352
+ if index:
353
+ idx = CategoricalIndex(idx)
354
+
355
+ result = ser[idx]
356
+ expected = ser[idx.fillna(False)]
357
+
358
+ tm.assert_series_equal(result, expected)
359
+
360
+
361
+ @pytest.fixture
362
+ def non_coercible_categorical(monkeypatch):
363
+ """
364
+ Monkeypatch Categorical.__array__ to ensure no implicit conversion.
365
+
366
+ Raises
367
+ ------
368
+ ValueError
369
+ When Categorical.__array__ is called.
370
+ """
371
+
372
+ # TODO(Categorical): identify other places where this may be
373
+ # useful and move to a conftest.py
374
+ def array(self, dtype=None):
375
+ raise ValueError("I cannot be converted.")
376
+
377
+ with monkeypatch.context() as m:
378
+ m.setattr(Categorical, "__array__", array)
379
+ yield
380
+
381
+
382
+ def test_series_at():
383
+ arr = Categorical(["a", "b", "c"])
384
+ ser = Series(arr)
385
+ result = ser.at[0]
386
+ assert result == "a"
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_map.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ from pandas import (
6
+ Categorical,
7
+ Index,
8
+ Series,
9
+ )
10
+ import pandas._testing as tm
11
+
12
+
13
+ @pytest.mark.parametrize(
14
+ "data, categories",
15
+ [
16
+ (list("abcbca"), list("cab")),
17
+ (pd.interval_range(0, 3).repeat(3), pd.interval_range(0, 3)),
18
+ ],
19
+ ids=["string", "interval"],
20
+ )
21
+ def test_map_str(data, categories, ordered, na_action):
22
+ # GH 31202 - override base class since we want to maintain categorical/ordered
23
+ cat = Categorical(data, categories=categories, ordered=ordered)
24
+ result = cat.map(str, na_action=na_action)
25
+ expected = Categorical(
26
+ map(str, data), categories=map(str, categories), ordered=ordered
27
+ )
28
+ tm.assert_categorical_equal(result, expected)
29
+
30
+
31
+ def test_map(na_action):
32
+ cat = Categorical(list("ABABC"), categories=list("CBA"), ordered=True)
33
+ result = cat.map(lambda x: x.lower(), na_action=na_action)
34
+ exp = Categorical(list("ababc"), categories=list("cba"), ordered=True)
35
+ tm.assert_categorical_equal(result, exp)
36
+
37
+ cat = Categorical(list("ABABC"), categories=list("BAC"), ordered=False)
38
+ result = cat.map(lambda x: x.lower(), na_action=na_action)
39
+ exp = Categorical(list("ababc"), categories=list("bac"), ordered=False)
40
+ tm.assert_categorical_equal(result, exp)
41
+
42
+ # GH 12766: Return an index not an array
43
+ result = cat.map(lambda x: 1, na_action=na_action)
44
+ exp = Index(np.array([1] * 5, dtype=np.int64))
45
+ tm.assert_index_equal(result, exp)
46
+
47
+ # change categories dtype
48
+ cat = Categorical(list("ABABC"), categories=list("BAC"), ordered=False)
49
+
50
+ def f(x):
51
+ return {"A": 10, "B": 20, "C": 30}.get(x)
52
+
53
+ result = cat.map(f, na_action=na_action)
54
+ exp = Categorical([10, 20, 10, 20, 30], categories=[20, 10, 30], ordered=False)
55
+ tm.assert_categorical_equal(result, exp)
56
+
57
+ mapper = Series([10, 20, 30], index=["A", "B", "C"])
58
+ result = cat.map(mapper, na_action=na_action)
59
+ tm.assert_categorical_equal(result, exp)
60
+
61
+ result = cat.map({"A": 10, "B": 20, "C": 30}, na_action=na_action)
62
+ tm.assert_categorical_equal(result, exp)
63
+
64
+
65
+ @pytest.mark.parametrize(
66
+ ("data", "f", "expected"),
67
+ (
68
+ ([1, 1, np.nan], pd.isna, Index([False, False, True])),
69
+ ([1, 2, np.nan], pd.isna, Index([False, False, True])),
70
+ ([1, 1, np.nan], {1: False}, Categorical([False, False, np.nan])),
71
+ ([1, 2, np.nan], {1: False, 2: False}, Index([False, False, np.nan])),
72
+ (
73
+ [1, 1, np.nan],
74
+ Series([False, False]),
75
+ Categorical([False, False, np.nan]),
76
+ ),
77
+ (
78
+ [1, 2, np.nan],
79
+ Series([False] * 3),
80
+ Index([False, False, np.nan]),
81
+ ),
82
+ ),
83
+ )
84
+ def test_map_with_nan_none(data, f, expected): # GH 24241
85
+ values = Categorical(data)
86
+ result = values.map(f, na_action=None)
87
+ if isinstance(expected, Categorical):
88
+ tm.assert_categorical_equal(result, expected)
89
+ else:
90
+ tm.assert_index_equal(result, expected)
91
+
92
+
93
+ @pytest.mark.parametrize(
94
+ ("data", "f", "expected"),
95
+ (
96
+ ([1, 1, np.nan], pd.isna, Categorical([False, False, np.nan])),
97
+ ([1, 2, np.nan], pd.isna, Index([False, False, np.nan])),
98
+ ([1, 1, np.nan], {1: False}, Categorical([False, False, np.nan])),
99
+ ([1, 2, np.nan], {1: False, 2: False}, Index([False, False, np.nan])),
100
+ (
101
+ [1, 1, np.nan],
102
+ Series([False, False]),
103
+ Categorical([False, False, np.nan]),
104
+ ),
105
+ (
106
+ [1, 2, np.nan],
107
+ Series([False, False, False]),
108
+ Index([False, False, np.nan]),
109
+ ),
110
+ ),
111
+ )
112
+ def test_map_with_nan_ignore(data, f, expected): # GH 24241
113
+ values = Categorical(data)
114
+ result = values.map(f, na_action="ignore")
115
+ if data[1] == 1:
116
+ tm.assert_categorical_equal(result, expected)
117
+ else:
118
+ tm.assert_index_equal(result, expected)
119
+
120
+
121
+ def test_map_with_dict_or_series(na_action):
122
+ orig_values = ["a", "B", 1, "a"]
123
+ new_values = ["one", 2, 3.0, "one"]
124
+ cat = Categorical(orig_values)
125
+
126
+ mapper = Series(new_values[:-1], index=orig_values[:-1])
127
+ result = cat.map(mapper, na_action=na_action)
128
+
129
+ # Order of categories in result can be different
130
+ expected = Categorical(new_values, categories=[3.0, 2, "one"])
131
+ tm.assert_categorical_equal(result, expected)
132
+
133
+ mapper = dict(zip(orig_values[:-1], new_values[:-1], strict=True))
134
+ result = cat.map(mapper, na_action=na_action)
135
+ # Order of categories in result can be different
136
+ tm.assert_categorical_equal(result, expected)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_missing.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import collections
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ from pandas.errors import Pandas4Warning
7
+
8
+ from pandas.core.dtypes.dtypes import CategoricalDtype
9
+
10
+ import pandas as pd
11
+ from pandas import (
12
+ Categorical,
13
+ Index,
14
+ Series,
15
+ isna,
16
+ )
17
+ import pandas._testing as tm
18
+
19
+
20
+ class TestCategoricalMissing:
21
+ def test_isna(self):
22
+ exp = np.array([False, False, True])
23
+ cat = Categorical(["a", "b", np.nan])
24
+ res = cat.isna()
25
+
26
+ tm.assert_numpy_array_equal(res, exp)
27
+
28
+ def test_na_flags_int_categories(self):
29
+ # #1457
30
+
31
+ categories = list(range(10))
32
+ labels = np.random.default_rng(2).integers(0, 10, 20)
33
+ labels[::5] = -1
34
+ msg = "Constructing a Categorical with a dtype and values containing"
35
+ with tm.assert_produces_warning(Pandas4Warning, match=msg):
36
+ cat = Categorical(labels, categories)
37
+ repr(cat)
38
+
39
+ tm.assert_numpy_array_equal(isna(cat), labels == -1)
40
+
41
+ def test_nan_handling(self):
42
+ # Nans are represented as -1 in codes
43
+ c = Categorical(["a", "b", np.nan, "a"])
44
+ tm.assert_index_equal(c.categories, Index(["a", "b"]))
45
+ tm.assert_numpy_array_equal(c._codes, np.array([0, 1, -1, 0], dtype=np.int8))
46
+ c[1] = np.nan
47
+ tm.assert_index_equal(c.categories, Index(["a", "b"]))
48
+ tm.assert_numpy_array_equal(c._codes, np.array([0, -1, -1, 0], dtype=np.int8))
49
+
50
+ # Adding nan to categories should make assigned nan point to the
51
+ # category!
52
+ c = Categorical(["a", "b", np.nan, "a"])
53
+ tm.assert_index_equal(c.categories, Index(["a", "b"]))
54
+ tm.assert_numpy_array_equal(c._codes, np.array([0, 1, -1, 0], dtype=np.int8))
55
+
56
+ def test_set_dtype_nans(self):
57
+ c = Categorical(["a", "b", np.nan])
58
+ result = c._set_dtype(CategoricalDtype(["a", "c"]), copy=True)
59
+ tm.assert_numpy_array_equal(result.codes, np.array([0, -1, -1], dtype="int8"))
60
+
61
+ def test_set_item_nan(self):
62
+ cat = Categorical([1, 2, 3])
63
+ cat[1] = np.nan
64
+
65
+ exp = Categorical([1, np.nan, 3], categories=[1, 2, 3])
66
+ tm.assert_categorical_equal(cat, exp)
67
+
68
+ @pytest.mark.parametrize("named", [True, False])
69
+ def test_fillna_iterable_category(self, named):
70
+ # https://github.com/pandas-dev/pandas/issues/21097
71
+ if named:
72
+ Point = collections.namedtuple("Point", "x y")
73
+ else:
74
+ Point = lambda *args: args # tuple
75
+ cat = Categorical(np.array([Point(0, 0), Point(0, 1), None], dtype=object))
76
+ result = cat.fillna(Point(0, 0))
77
+ expected = Categorical([Point(0, 0), Point(0, 1), Point(0, 0)])
78
+
79
+ tm.assert_categorical_equal(result, expected)
80
+
81
+ # Case where the Point is not among our categories; we want ValueError,
82
+ # not NotImplementedError GH#41914
83
+ cat = Categorical(np.array([Point(1, 0), Point(0, 1), None], dtype=object))
84
+ msg = "Cannot setitem on a Categorical with a new category"
85
+ with pytest.raises(TypeError, match=msg):
86
+ cat.fillna(Point(0, 0))
87
+
88
+ def test_fillna_array(self):
89
+ # accept Categorical or ndarray value if it holds appropriate values
90
+ cat = Categorical(["A", "B", "C", None, None])
91
+
92
+ other = cat.fillna("C")
93
+ result = cat.fillna(other)
94
+ tm.assert_categorical_equal(result, other)
95
+ assert isna(cat[-1]) # didn't modify original inplace
96
+
97
+ other = np.array(["A", "B", "C", "B", "A"])
98
+ result = cat.fillna(other)
99
+ expected = Categorical(["A", "B", "C", "B", "A"], dtype=cat.dtype)
100
+ tm.assert_categorical_equal(result, expected)
101
+ assert isna(cat[-1]) # didn't modify original inplace
102
+
103
+ @pytest.mark.parametrize(
104
+ "a1, a2, categories",
105
+ [
106
+ (["a", "b", "c"], [np.nan, "a", "b"], ["a", "b", "c"]),
107
+ ([1, 2, 3], [np.nan, 1, 2], [1, 2, 3]),
108
+ ],
109
+ )
110
+ def test_compare_categorical_with_missing(self, a1, a2, categories):
111
+ # GH 28384
112
+ cat_type = CategoricalDtype(categories)
113
+
114
+ # !=
115
+ result = Series(a1, dtype=cat_type) != Series(a2, dtype=cat_type)
116
+ expected = Series(a1) != Series(a2)
117
+ tm.assert_series_equal(result, expected)
118
+
119
+ # ==
120
+ result = Series(a1, dtype=cat_type) == Series(a2, dtype=cat_type)
121
+ expected = Series(a1) == Series(a2)
122
+ tm.assert_series_equal(result, expected)
123
+
124
+ @pytest.mark.parametrize(
125
+ "na_value, dtype",
126
+ [
127
+ (pd.NaT, "datetime64[s]"),
128
+ (None, "float64"),
129
+ (np.nan, "float64"),
130
+ (pd.NA, "float64"),
131
+ ],
132
+ )
133
+ def test_categorical_only_missing_values_no_cast(self, na_value, dtype):
134
+ # GH#44900
135
+ result = Categorical([na_value, na_value])
136
+ tm.assert_index_equal(result.categories, Index([], dtype=dtype))
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_operators.py ADDED
@@ -0,0 +1,408 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ import pandas as pd
5
+ from pandas import (
6
+ Categorical,
7
+ DataFrame,
8
+ Series,
9
+ Timestamp,
10
+ date_range,
11
+ )
12
+ import pandas._testing as tm
13
+
14
+
15
+ class TestCategoricalOpsWithFactor:
16
+ def test_categories_none_comparisons(self):
17
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
18
+ tm.assert_categorical_equal(factor, factor)
19
+
20
+ def test_comparisons(self):
21
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
22
+ result = factor[factor == "a"]
23
+ expected = factor[np.asarray(factor) == "a"]
24
+ tm.assert_categorical_equal(result, expected)
25
+
26
+ result = factor[factor != "a"]
27
+ expected = factor[np.asarray(factor) != "a"]
28
+ tm.assert_categorical_equal(result, expected)
29
+
30
+ result = factor[factor < "c"]
31
+ expected = factor[np.asarray(factor) < "c"]
32
+ tm.assert_categorical_equal(result, expected)
33
+
34
+ result = factor[factor > "a"]
35
+ expected = factor[np.asarray(factor) > "a"]
36
+ tm.assert_categorical_equal(result, expected)
37
+
38
+ result = factor[factor >= "b"]
39
+ expected = factor[np.asarray(factor) >= "b"]
40
+ tm.assert_categorical_equal(result, expected)
41
+
42
+ result = factor[factor <= "b"]
43
+ expected = factor[np.asarray(factor) <= "b"]
44
+ tm.assert_categorical_equal(result, expected)
45
+
46
+ n = len(factor)
47
+
48
+ other = factor[np.random.default_rng(2).permutation(n)]
49
+ result = factor == other
50
+ expected = np.asarray(factor) == np.asarray(other)
51
+ tm.assert_numpy_array_equal(result, expected)
52
+
53
+ result = factor == "d"
54
+ expected = np.zeros(len(factor), dtype=bool)
55
+ tm.assert_numpy_array_equal(result, expected)
56
+
57
+ # comparisons with categoricals
58
+ cat_rev = Categorical(["a", "b", "c"], categories=["c", "b", "a"], ordered=True)
59
+ cat_rev_base = Categorical(
60
+ ["b", "b", "b"], categories=["c", "b", "a"], ordered=True
61
+ )
62
+ cat = Categorical(["a", "b", "c"], ordered=True)
63
+ cat_base = Categorical(["b", "b", "b"], categories=cat.categories, ordered=True)
64
+
65
+ # comparisons need to take categories ordering into account
66
+ res_rev = cat_rev > cat_rev_base
67
+ exp_rev = np.array([True, False, False])
68
+ tm.assert_numpy_array_equal(res_rev, exp_rev)
69
+
70
+ res_rev = cat_rev < cat_rev_base
71
+ exp_rev = np.array([False, False, True])
72
+ tm.assert_numpy_array_equal(res_rev, exp_rev)
73
+
74
+ res = cat > cat_base
75
+ exp = np.array([False, False, True])
76
+ tm.assert_numpy_array_equal(res, exp)
77
+
78
+ # Only categories with same categories can be compared
79
+ msg = "Categoricals can only be compared if 'categories' are the same"
80
+ with pytest.raises(TypeError, match=msg):
81
+ cat > cat_rev
82
+
83
+ cat_rev_base2 = Categorical(["b", "b", "b"], categories=["c", "b", "a", "d"])
84
+
85
+ with pytest.raises(TypeError, match=msg):
86
+ cat_rev > cat_rev_base2
87
+
88
+ # Only categories with same ordering information can be compared
89
+ cat_unordered = cat.set_ordered(False)
90
+ assert not (cat > cat).any()
91
+
92
+ with pytest.raises(TypeError, match=msg):
93
+ cat > cat_unordered
94
+
95
+ # comparison (in both directions) with Series will raise
96
+ s = Series(["b", "b", "b"], dtype=object)
97
+ msg = (
98
+ "Cannot compare a Categorical for op __gt__ with type "
99
+ r"<class 'numpy\.ndarray'>"
100
+ )
101
+ with pytest.raises(TypeError, match=msg):
102
+ cat > s
103
+ with pytest.raises(TypeError, match=msg):
104
+ cat_rev > s
105
+ with pytest.raises(TypeError, match=msg):
106
+ s < cat
107
+ with pytest.raises(TypeError, match=msg):
108
+ s < cat_rev
109
+
110
+ # comparison with numpy.array will raise in both direction, but only on
111
+ # newer numpy versions
112
+ a = np.array(["b", "b", "b"], dtype=object)
113
+ with pytest.raises(TypeError, match=msg):
114
+ cat > a
115
+ with pytest.raises(TypeError, match=msg):
116
+ cat_rev > a
117
+
118
+ # Make sure that unequal comparison take the categories order in
119
+ # account
120
+ cat_rev = Categorical(list("abc"), categories=list("cba"), ordered=True)
121
+ exp = np.array([True, False, False])
122
+ res = cat_rev > "b"
123
+ tm.assert_numpy_array_equal(res, exp)
124
+
125
+ # check that zero-dim array gets unboxed
126
+ res = cat_rev > np.array("b")
127
+ tm.assert_numpy_array_equal(res, exp)
128
+
129
+
130
+ class TestCategoricalOps:
131
+ @pytest.mark.parametrize(
132
+ "categories",
133
+ [["a", "b"], [0, 1], [Timestamp("2019"), Timestamp("2020")]],
134
+ )
135
+ def test_not_equal_with_na(self, categories):
136
+ # https://github.com/pandas-dev/pandas/issues/32276
137
+ c1 = Categorical.from_codes([-1, 0], categories=categories)
138
+ c2 = Categorical.from_codes([0, 1], categories=categories)
139
+
140
+ result = c1 != c2
141
+
142
+ assert result.all()
143
+
144
+ def test_compare_frame(self):
145
+ # GH#24282 check that Categorical.__cmp__(DataFrame) defers to frame
146
+ data = ["a", "b", 2, "a"]
147
+ cat = Categorical(data)
148
+
149
+ df = DataFrame(cat)
150
+
151
+ result = cat == df.T
152
+ expected = DataFrame([[True, True, True, True]])
153
+ tm.assert_frame_equal(result, expected)
154
+
155
+ result = cat[::-1] != df.T
156
+ expected = DataFrame([[False, True, True, False]])
157
+ tm.assert_frame_equal(result, expected)
158
+
159
+ def test_compare_frame_raises(self, comparison_op):
160
+ # alignment raises unless we transpose
161
+ op = comparison_op
162
+ cat = Categorical(["a", "b", 2, "a"])
163
+ df = DataFrame(cat)
164
+ msg = "Unable to coerce to Series, length must be 1: given 4"
165
+ with pytest.raises(ValueError, match=msg):
166
+ op(cat, df)
167
+
168
+ def test_datetime_categorical_comparison(self):
169
+ dt_cat = Categorical(date_range("2014-01-01", periods=3), ordered=True)
170
+ tm.assert_numpy_array_equal(dt_cat > dt_cat[0], np.array([False, True, True]))
171
+ tm.assert_numpy_array_equal(dt_cat[0] < dt_cat, np.array([False, True, True]))
172
+
173
+ def test_reflected_comparison_with_scalars(self):
174
+ # GH8658
175
+ cat = Categorical([1, 2, 3], ordered=True)
176
+ tm.assert_numpy_array_equal(cat > cat[0], np.array([False, True, True]))
177
+ tm.assert_numpy_array_equal(cat[0] < cat, np.array([False, True, True]))
178
+
179
+ def test_comparison_with_unknown_scalars(self):
180
+ # https://github.com/pandas-dev/pandas/issues/9836#issuecomment-92123057
181
+ # and following comparisons with scalars not in categories should raise
182
+ # for unequal comps, but not for equal/not equal
183
+ cat = Categorical([1, 2, 3], ordered=True)
184
+
185
+ msg = "Invalid comparison between dtype=category and int"
186
+ with pytest.raises(TypeError, match=msg):
187
+ cat < 4
188
+ with pytest.raises(TypeError, match=msg):
189
+ cat > 4
190
+ with pytest.raises(TypeError, match=msg):
191
+ 4 < cat
192
+ with pytest.raises(TypeError, match=msg):
193
+ 4 > cat
194
+
195
+ tm.assert_numpy_array_equal(cat == 4, np.array([False, False, False]))
196
+ tm.assert_numpy_array_equal(cat != 4, np.array([True, True, True]))
197
+
198
+ def test_comparison_with_tuple(self):
199
+ cat = Categorical(np.array(["foo", (0, 1), 3, (0, 1)], dtype=object))
200
+
201
+ result = cat == "foo"
202
+ expected = np.array([True, False, False, False], dtype=bool)
203
+ tm.assert_numpy_array_equal(result, expected)
204
+
205
+ result = cat == (0, 1)
206
+ expected = np.array([False, True, False, True], dtype=bool)
207
+ tm.assert_numpy_array_equal(result, expected)
208
+
209
+ result = cat != (0, 1)
210
+ tm.assert_numpy_array_equal(result, ~expected)
211
+
212
+ @pytest.mark.filterwarnings("ignore::RuntimeWarning")
213
+ def test_comparison_of_ordered_categorical_with_nan_to_scalar(
214
+ self, compare_operators_no_eq_ne
215
+ ):
216
+ # https://github.com/pandas-dev/pandas/issues/26504
217
+ # BUG: fix ordered categorical comparison with missing values (#26504 )
218
+ # and following comparisons with scalars in categories with missing
219
+ # values should be evaluated as False
220
+
221
+ cat = Categorical([1, 2, 3, None], categories=[1, 2, 3], ordered=True)
222
+ scalar = 2
223
+ expected = getattr(np.array(cat), compare_operators_no_eq_ne)(scalar)
224
+ actual = getattr(cat, compare_operators_no_eq_ne)(scalar)
225
+ tm.assert_numpy_array_equal(actual, expected)
226
+
227
+ @pytest.mark.filterwarnings("ignore::RuntimeWarning")
228
+ def test_comparison_of_ordered_categorical_with_nan_to_listlike(
229
+ self, compare_operators_no_eq_ne
230
+ ):
231
+ # https://github.com/pandas-dev/pandas/issues/26504
232
+ # and following comparisons of missing values in ordered Categorical
233
+ # with listlike should be evaluated as False
234
+
235
+ cat = Categorical([1, 2, 3, None], categories=[1, 2, 3], ordered=True)
236
+ other = Categorical([2, 2, 2, 2], categories=[1, 2, 3], ordered=True)
237
+ expected = getattr(np.array(cat), compare_operators_no_eq_ne)(2)
238
+ actual = getattr(cat, compare_operators_no_eq_ne)(other)
239
+ tm.assert_numpy_array_equal(actual, expected)
240
+
241
+ @pytest.mark.parametrize(
242
+ "data,reverse,base",
243
+ [(list("abc"), list("cba"), list("bbb")), ([1, 2, 3], [3, 2, 1], [2, 2, 2])],
244
+ )
245
+ def test_comparisons(self, data, reverse, base):
246
+ cat_rev = Series(Categorical(data, categories=reverse, ordered=True))
247
+ cat_rev_base = Series(Categorical(base, categories=reverse, ordered=True))
248
+ cat = Series(Categorical(data, ordered=True))
249
+ cat_base = Series(
250
+ Categorical(base, categories=cat.cat.categories, ordered=True)
251
+ )
252
+ s = Series(base, dtype=object if base == list("bbb") else None)
253
+ a = np.array(base)
254
+
255
+ # comparisons need to take categories ordering into account
256
+ res_rev = cat_rev > cat_rev_base
257
+ exp_rev = Series([True, False, False])
258
+ tm.assert_series_equal(res_rev, exp_rev)
259
+
260
+ res_rev = cat_rev < cat_rev_base
261
+ exp_rev = Series([False, False, True])
262
+ tm.assert_series_equal(res_rev, exp_rev)
263
+
264
+ res = cat > cat_base
265
+ exp = Series([False, False, True])
266
+ tm.assert_series_equal(res, exp)
267
+
268
+ scalar = base[1]
269
+ res = cat > scalar
270
+ exp = Series([False, False, True])
271
+ exp2 = cat.values > scalar
272
+ tm.assert_series_equal(res, exp)
273
+ tm.assert_numpy_array_equal(res.values, exp2)
274
+ res_rev = cat_rev > scalar
275
+ exp_rev = Series([True, False, False])
276
+ exp_rev2 = cat_rev.values > scalar
277
+ tm.assert_series_equal(res_rev, exp_rev)
278
+ tm.assert_numpy_array_equal(res_rev.values, exp_rev2)
279
+
280
+ # Only categories with same categories can be compared
281
+ msg = "Categoricals can only be compared if 'categories' are the same"
282
+ with pytest.raises(TypeError, match=msg):
283
+ cat > cat_rev
284
+
285
+ # categorical cannot be compared to Series or numpy array, and also
286
+ # not the other way around
287
+ msg = (
288
+ "Cannot compare a Categorical for op __gt__ with type "
289
+ r"<class 'numpy\.ndarray'>"
290
+ )
291
+ with pytest.raises(TypeError, match=msg):
292
+ cat > s
293
+ with pytest.raises(TypeError, match=msg):
294
+ cat_rev > s
295
+ with pytest.raises(TypeError, match=msg):
296
+ cat > a
297
+ with pytest.raises(TypeError, match=msg):
298
+ cat_rev > a
299
+
300
+ with pytest.raises(TypeError, match=msg):
301
+ s < cat
302
+ with pytest.raises(TypeError, match=msg):
303
+ s < cat_rev
304
+
305
+ with pytest.raises(TypeError, match=msg):
306
+ a < cat
307
+ with pytest.raises(TypeError, match=msg):
308
+ a < cat_rev
309
+
310
+ @pytest.mark.parametrize("box", [lambda x: x, Series])
311
+ def test_unordered_different_order_equal(self, box):
312
+ # https://github.com/pandas-dev/pandas/issues/16014
313
+ c1 = box(Categorical(["a", "b"], categories=["a", "b"], ordered=False))
314
+ c2 = box(Categorical(["a", "b"], categories=["b", "a"], ordered=False))
315
+ assert (c1 == c2).all()
316
+
317
+ c1 = box(Categorical(["a", "b"], categories=["a", "b"], ordered=False))
318
+ c2 = box(Categorical(["b", "a"], categories=["b", "a"], ordered=False))
319
+ assert (c1 != c2).all()
320
+
321
+ c1 = box(Categorical(["a", "a"], categories=["a", "b"], ordered=False))
322
+ c2 = box(Categorical(["b", "b"], categories=["b", "a"], ordered=False))
323
+ assert (c1 != c2).all()
324
+
325
+ c1 = box(Categorical(["a", "a"], categories=["a", "b"], ordered=False))
326
+ c2 = box(Categorical(["a", "b"], categories=["b", "a"], ordered=False))
327
+ result = c1 == c2
328
+ tm.assert_numpy_array_equal(np.array(result), np.array([True, False]))
329
+
330
+ def test_unordered_different_categories_raises(self):
331
+ c1 = Categorical(["a", "b"], categories=["a", "b"], ordered=False)
332
+ c2 = Categorical(["a", "c"], categories=["c", "a"], ordered=False)
333
+
334
+ with pytest.raises(TypeError, match=("Categoricals can only be compared")):
335
+ c1 == c2
336
+
337
+ def test_compare_different_lengths(self):
338
+ c1 = Categorical([], categories=["a", "b"])
339
+ c2 = Categorical([], categories=["a"])
340
+
341
+ msg = "Categoricals can only be compared if 'categories' are the same."
342
+ with pytest.raises(TypeError, match=msg):
343
+ c1 == c2
344
+
345
+ def test_compare_unordered_different_order(self):
346
+ # https://github.com/pandas-dev/pandas/issues/16603#issuecomment-
347
+ # 349290078
348
+ a = Categorical(["a"], categories=["a", "b"])
349
+ b = Categorical(["b"], categories=["b", "a"])
350
+ assert not a.equals(b)
351
+
352
+ def test_numeric_like_ops(self):
353
+ df = DataFrame({"value": np.random.default_rng(2).integers(0, 10000, 100)})
354
+ labels = [f"{i} - {i + 499}" for i in range(0, 10000, 500)]
355
+ cat_labels = Categorical(labels, labels)
356
+
357
+ df = df.sort_values(by=["value"], ascending=True)
358
+ df["value_group"] = pd.cut(
359
+ df.value, range(0, 10500, 500), right=False, labels=cat_labels
360
+ )
361
+
362
+ # numeric ops should not succeed
363
+ for op, str_rep in [
364
+ ("__add__", r"\+"),
365
+ ("__sub__", "-"),
366
+ ("__mul__", r"\*"),
367
+ ("__truediv__", "/"),
368
+ ]:
369
+ msg = f"Series cannot perform the operation {str_rep}|unsupported operand"
370
+ with pytest.raises(TypeError, match=msg):
371
+ getattr(df, op)(df)
372
+
373
+ # reduction ops should not succeed (unless specifically defined, e.g.
374
+ # min/max)
375
+ s = df["value_group"]
376
+ for op in ["kurt", "skew", "var", "std", "mean", "sum", "median"]:
377
+ msg = f"does not support operation '{op}'"
378
+ with pytest.raises(TypeError, match=msg):
379
+ getattr(s, op)(numeric_only=False)
380
+
381
+ def test_numeric_like_ops_series(self):
382
+ # numpy ops
383
+ s = Series(Categorical([1, 2, 3, 4]))
384
+ with pytest.raises(TypeError, match="does not support operation 'sum'"):
385
+ np.sum(s)
386
+
387
+ @pytest.mark.parametrize(
388
+ "op, str_rep",
389
+ [
390
+ ("__add__", r"\+"),
391
+ ("__sub__", "-"),
392
+ ("__mul__", r"\*"),
393
+ ("__truediv__", "/"),
394
+ ],
395
+ )
396
+ def test_numeric_like_ops_series_arith(self, op, str_rep):
397
+ # numeric ops on a Series
398
+ s = Series(Categorical([1, 2, 3, 4]))
399
+ msg = f"Series cannot perform the operation {str_rep}|unsupported operand"
400
+ with pytest.raises(TypeError, match=msg):
401
+ getattr(s, op)(2)
402
+
403
+ def test_numeric_like_ops_series_invalid(self):
404
+ # invalid ufunc
405
+ s = Series(Categorical([1, 2, 3, 4]))
406
+ msg = "Object with dtype category cannot perform the numpy op log"
407
+ with pytest.raises(TypeError, match=msg):
408
+ np.log(s)
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_replace.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pytest
2
+
3
+ import pandas as pd
4
+ from pandas import Categorical
5
+ import pandas._testing as tm
6
+
7
+
8
+ @pytest.mark.parametrize(
9
+ "to_replace,value,expected",
10
+ [
11
+ # one-to-one
12
+ (4, 1, [1, 2, 3]),
13
+ (3, 1, [1, 2, 1]),
14
+ # many-to-one
15
+ ((5, 6), 2, [1, 2, 3]),
16
+ ((3, 2), 1, [1, 1, 1]),
17
+ ],
18
+ )
19
+ def test_replace_categorical_series(to_replace, value, expected):
20
+ # GH 31720
21
+ ser = pd.Series([1, 2, 3], dtype="category")
22
+ result = ser.replace(to_replace, value)
23
+ expected = pd.Series(Categorical(expected, categories=[1, 2, 3]))
24
+ tm.assert_series_equal(result, expected)
25
+
26
+
27
+ @pytest.mark.parametrize(
28
+ "to_replace,value",
29
+ [
30
+ # one-to-one
31
+ (3, 5),
32
+ # many-to-one
33
+ ((3, 2), 5),
34
+ ],
35
+ )
36
+ def test_replace_categorical_series_new_category_raises(to_replace, value):
37
+ # GH 31720
38
+ ser = pd.Series([1, 2, 3], dtype="category")
39
+ with pytest.raises(
40
+ TypeError, match="Cannot setitem on a Categorical with a new category"
41
+ ):
42
+ ser.replace(to_replace, value)
43
+
44
+
45
+ def test_replace_maintain_ordering():
46
+ # GH51016
47
+ dtype = pd.CategoricalDtype([0, 1, 2], ordered=True)
48
+ ser = pd.Series([0, 1, 2], dtype=dtype)
49
+ result = ser.replace(0, 2)
50
+ expected = pd.Series([2, 1, 2], dtype=dtype)
51
+ tm.assert_series_equal(expected, result, check_category_order=True)
52
+
53
+
54
+ def test_replace_categorical_ea_dtype():
55
+ # GH49404
56
+ cat = Categorical(pd.array(["a", "b", "c"], dtype="string"))
57
+ result = pd.Series(cat).replace(["a", "b"], ["c", "c"])._values
58
+ expected = Categorical(
59
+ pd.array(["c"] * 3, dtype="string"),
60
+ categories=pd.array(["a", "b", "c"], dtype="string"),
61
+ )
62
+ tm.assert_categorical_equal(result, expected)
63
+
64
+
65
+ def test_replace_categorical_ea_dtype_different_cats_raises():
66
+ # GH49404
67
+ cat = Categorical(pd.array(["a", "b"], dtype="string"))
68
+ with pytest.raises(
69
+ TypeError, match="Cannot setitem on a Categorical with a new category"
70
+ ):
71
+ pd.Series(cat).replace(["a", "b"], ["c", pd.NA])
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_repr.py ADDED
@@ -0,0 +1,561 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ from pandas import (
4
+ Categorical,
5
+ CategoricalDtype,
6
+ CategoricalIndex,
7
+ Index,
8
+ Series,
9
+ date_range,
10
+ option_context,
11
+ period_range,
12
+ timedelta_range,
13
+ )
14
+
15
+
16
+ class TestCategoricalReprWithFactor:
17
+ def test_print(self, using_infer_string):
18
+ factor = Categorical(["a", "b", "b", "a", "a", "c", "c", "c"], ordered=True)
19
+ dtype = "str" if using_infer_string else "object"
20
+ expected = [
21
+ "['a', 'b', 'b', 'a', 'a', 'c', 'c', 'c']",
22
+ f"Categories (3, {dtype}): ['a' < 'b' < 'c']",
23
+ ]
24
+ expected = "\n".join(expected)
25
+ actual = repr(factor)
26
+ assert actual == expected
27
+
28
+
29
+ class TestCategoricalRepr:
30
+ def test_big_print(self):
31
+ codes = np.array([0, 1, 2, 0, 1, 2] * 100)
32
+ dtype = CategoricalDtype(categories=Index(["a", "b", "c"], dtype=object))
33
+ factor = Categorical.from_codes(codes, dtype=dtype)
34
+ expected = [
35
+ "['a', 'b', 'c', 'a', 'b', ..., 'b', 'c', 'a', 'b', 'c']",
36
+ "Length: 600",
37
+ "Categories (3, object): ['a', 'b', 'c']",
38
+ ]
39
+ expected = "\n".join(expected)
40
+
41
+ actual = repr(factor)
42
+
43
+ assert actual == expected
44
+
45
+ def test_empty_print(self):
46
+ factor = Categorical([], Index(["a", "b", "c"], dtype=object))
47
+ expected = "[], Categories (3, object): ['a', 'b', 'c']"
48
+ actual = repr(factor)
49
+ assert actual == expected
50
+
51
+ assert expected == actual
52
+ factor = Categorical([], Index(["a", "b", "c"], dtype=object), ordered=True)
53
+ expected = "[], Categories (3, object): ['a' < 'b' < 'c']"
54
+ actual = repr(factor)
55
+ assert expected == actual
56
+
57
+ factor = Categorical([], [])
58
+ expected = "[], Categories (0, object): []"
59
+ assert expected == repr(factor)
60
+
61
+ def test_print_none_width(self):
62
+ # GH10087
63
+ a = Series(Categorical([1, 2, 3, 4]))
64
+ exp = (
65
+ "0 1\n1 2\n2 3\n3 4\n"
66
+ "dtype: category\nCategories (4, int64): [1, 2, 3, 4]"
67
+ )
68
+
69
+ with option_context("display.width", None):
70
+ assert exp == repr(a)
71
+
72
+ def test_unicode_print(self, using_infer_string):
73
+ c = Categorical(["aaaaa", "bb", "cccc"] * 20)
74
+ expected = """\
75
+ ['aaaaa', 'bb', 'cccc', 'aaaaa', 'bb', ..., 'bb', 'cccc', 'aaaaa', 'bb', 'cccc']
76
+ Length: 60
77
+ Categories (3, object): ['aaaaa', 'bb', 'cccc']"""
78
+
79
+ if using_infer_string:
80
+ expected = expected.replace("object", "str")
81
+
82
+ assert repr(c) == expected
83
+
84
+ c = Categorical(["ああああ", "いいいいい", "ううううううう"] * 20)
85
+ expected = """\
86
+ ['ああああ', 'いいいいい', 'ううううううう', 'ああああ', 'いいいいい', ..., 'いいいいい', 'ううううううう', 'ああああ', 'いいいいい', 'ううううううう']
87
+ Length: 60
88
+ Categories (3, object): ['ああああ', 'いいいいい', 'ううううううう']""" # noqa: E501
89
+
90
+ if using_infer_string:
91
+ expected = expected.replace("object", "str")
92
+
93
+ assert repr(c) == expected
94
+
95
+ # unicode option should not affect to Categorical, as it doesn't care
96
+ # the repr width
97
+ with option_context("display.unicode.east_asian_width", True):
98
+ c = Categorical(["ああああ", "いいいいい", "ううううううう"] * 20)
99
+ expected = """['ああああ', 'いいいいい', 'ううううううう', 'ああああ', 'いいいいい', ..., 'いいいいい', 'ううううううう', 'ああああ', 'いいいいい', 'ううううううう']
100
+ Length: 60
101
+ Categories (3, object): ['ああああ', 'いいいいい', 'ううううううう']""" # noqa: E501
102
+
103
+ if using_infer_string:
104
+ expected = expected.replace("object", "str")
105
+
106
+ assert repr(c) == expected
107
+
108
+ def test_categorical_repr(self):
109
+ c = Categorical([1, 2, 3])
110
+ exp = """[1, 2, 3]
111
+ Categories (3, int64): [1, 2, 3]"""
112
+
113
+ assert repr(c) == exp
114
+
115
+ c = Categorical([1, 2, 3, 1, 2, 3], categories=[1, 2, 3])
116
+ exp = """[1, 2, 3, 1, 2, 3]
117
+ Categories (3, int64): [1, 2, 3]"""
118
+
119
+ assert repr(c) == exp
120
+
121
+ c = Categorical([1, 2, 3, 4, 5] * 10)
122
+ exp = """[1, 2, 3, 4, 5, ..., 1, 2, 3, 4, 5]
123
+ Length: 50
124
+ Categories (5, int64): [1, 2, 3, 4, 5]"""
125
+
126
+ assert repr(c) == exp
127
+
128
+ c = Categorical(np.arange(20, dtype=np.int64))
129
+ exp = """[0, 1, 2, 3, 4, ..., 15, 16, 17, 18, 19]
130
+ Length: 20
131
+ Categories (20, int64): [0, 1, 2, 3, ..., 16, 17, 18, 19]"""
132
+
133
+ assert repr(c) == exp
134
+
135
+ def test_categorical_repr_ordered(self):
136
+ c = Categorical([1, 2, 3], ordered=True)
137
+ exp = """[1, 2, 3]
138
+ Categories (3, int64): [1 < 2 < 3]"""
139
+
140
+ assert repr(c) == exp
141
+
142
+ c = Categorical([1, 2, 3, 1, 2, 3], categories=[1, 2, 3], ordered=True)
143
+ exp = """[1, 2, 3, 1, 2, 3]
144
+ Categories (3, int64): [1 < 2 < 3]"""
145
+
146
+ assert repr(c) == exp
147
+
148
+ c = Categorical([1, 2, 3, 4, 5] * 10, ordered=True)
149
+ exp = """[1, 2, 3, 4, 5, ..., 1, 2, 3, 4, 5]
150
+ Length: 50
151
+ Categories (5, int64): [1 < 2 < 3 < 4 < 5]"""
152
+
153
+ assert repr(c) == exp
154
+
155
+ c = Categorical(np.arange(20, dtype=np.int64), ordered=True)
156
+ exp = """[0, 1, 2, 3, 4, ..., 15, 16, 17, 18, 19]
157
+ Length: 20
158
+ Categories (20, int64): [0 < 1 < 2 < 3 ... 16 < 17 < 18 < 19]"""
159
+
160
+ assert repr(c) == exp
161
+
162
+ def test_categorical_repr_datetime(self):
163
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5, unit="ns")
164
+ c = Categorical(idx)
165
+
166
+ exp = (
167
+ "[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, "
168
+ "2011-01-01 12:00:00, 2011-01-01 13:00:00]\n"
169
+ "Categories (5, datetime64[ns]): [2011-01-01 09:00:00, "
170
+ "2011-01-01 10:00:00, 2011-01-01 11:00:00,\n"
171
+ " 2011-01-01 12:00:00, "
172
+ "2011-01-01 13:00:00]"
173
+ ""
174
+ )
175
+ assert repr(c) == exp
176
+
177
+ c = Categorical(idx.append(idx), categories=idx)
178
+ exp = (
179
+ "[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, "
180
+ "2011-01-01 12:00:00, 2011-01-01 13:00:00, 2011-01-01 09:00:00, "
181
+ "2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, "
182
+ "2011-01-01 13:00:00]\n"
183
+ "Categories (5, datetime64[ns]): [2011-01-01 09:00:00, "
184
+ "2011-01-01 10:00:00, 2011-01-01 11:00:00,\n"
185
+ " 2011-01-01 12:00:00, "
186
+ "2011-01-01 13:00:00]"
187
+ )
188
+
189
+ assert repr(c) == exp
190
+
191
+ idx = date_range(
192
+ "2011-01-01 09:00", freq="h", periods=5, tz="US/Eastern", unit="ns"
193
+ )
194
+ c = Categorical(idx)
195
+ exp = (
196
+ "[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, "
197
+ "2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, "
198
+ "2011-01-01 13:00:00-05:00]\n"
199
+ "Categories (5, datetime64[ns, US/Eastern]): "
200
+ "[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00,\n"
201
+ " "
202
+ "2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00,\n"
203
+ " "
204
+ "2011-01-01 13:00:00-05:00]"
205
+ )
206
+
207
+ assert repr(c) == exp
208
+
209
+ c = Categorical(idx.append(idx), categories=idx)
210
+ exp = (
211
+ "[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, "
212
+ "2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, "
213
+ "2011-01-01 13:00:00-05:00, 2011-01-01 09:00:00-05:00, "
214
+ "2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, "
215
+ "2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00]\n"
216
+ "Categories (5, datetime64[ns, US/Eastern]): "
217
+ "[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00,\n"
218
+ " "
219
+ "2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00,\n"
220
+ " "
221
+ "2011-01-01 13:00:00-05:00]"
222
+ )
223
+
224
+ assert repr(c) == exp
225
+
226
+ def test_categorical_repr_datetime_ordered(self):
227
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5, unit="ns")
228
+ c = Categorical(idx, ordered=True)
229
+ exp = """[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, 2011-01-01 13:00:00]
230
+ Categories (5, datetime64[ns]): [2011-01-01 09:00:00 < 2011-01-01 10:00:00 < 2011-01-01 11:00:00 <
231
+ 2011-01-01 12:00:00 < 2011-01-01 13:00:00]""" # noqa: E501
232
+
233
+ assert repr(c) == exp
234
+
235
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
236
+ exp = """[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, 2011-01-01 13:00:00, 2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, 2011-01-01 13:00:00]
237
+ Categories (5, datetime64[ns]): [2011-01-01 09:00:00 < 2011-01-01 10:00:00 < 2011-01-01 11:00:00 <
238
+ 2011-01-01 12:00:00 < 2011-01-01 13:00:00]""" # noqa: E501
239
+
240
+ assert repr(c) == exp
241
+
242
+ idx = date_range(
243
+ "2011-01-01 09:00", freq="h", periods=5, tz="US/Eastern", unit="ns"
244
+ )
245
+ c = Categorical(idx, ordered=True)
246
+ exp = """[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00]
247
+ Categories (5, datetime64[ns, US/Eastern]): [2011-01-01 09:00:00-05:00 < 2011-01-01 10:00:00-05:00 <
248
+ 2011-01-01 11:00:00-05:00 < 2011-01-01 12:00:00-05:00 <
249
+ 2011-01-01 13:00:00-05:00]""" # noqa: E501
250
+
251
+ assert repr(c) == exp
252
+
253
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
254
+ exp = """[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00, 2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00]
255
+ Categories (5, datetime64[ns, US/Eastern]): [2011-01-01 09:00:00-05:00 < 2011-01-01 10:00:00-05:00 <
256
+ 2011-01-01 11:00:00-05:00 < 2011-01-01 12:00:00-05:00 <
257
+ 2011-01-01 13:00:00-05:00]""" # noqa: E501
258
+
259
+ assert repr(c) == exp
260
+
261
+ def test_categorical_repr_int_with_nan(self):
262
+ c = Categorical([1, 2, np.nan])
263
+ c_exp = """[1, 2, NaN]\nCategories (2, int64): [1, 2]"""
264
+ assert repr(c) == c_exp
265
+
266
+ s = Series([1, 2, np.nan], dtype="object").astype("category")
267
+ s_exp = """0 1\n1 2\n2 NaN
268
+ dtype: category
269
+ Categories (2, int64): [1, 2]"""
270
+ assert repr(s) == s_exp
271
+
272
+ def test_categorical_repr_period(self):
273
+ idx = period_range("2011-01-01 09:00", freq="h", periods=5)
274
+ c = Categorical(idx)
275
+ exp = """[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00]
276
+ Categories (5, period[h]): [2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00,
277
+ 2011-01-01 13:00]""" # noqa: E501
278
+
279
+ assert repr(c) == exp
280
+
281
+ c = Categorical(idx.append(idx), categories=idx)
282
+ exp = """[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00, 2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00]
283
+ Categories (5, period[h]): [2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00,
284
+ 2011-01-01 13:00]""" # noqa: E501
285
+
286
+ assert repr(c) == exp
287
+
288
+ idx = period_range("2011-01", freq="M", periods=5)
289
+ c = Categorical(idx)
290
+ exp = """[2011-01, 2011-02, 2011-03, 2011-04, 2011-05]
291
+ Categories (5, period[M]): [2011-01, 2011-02, 2011-03, 2011-04, 2011-05]"""
292
+
293
+ assert repr(c) == exp
294
+
295
+ c = Categorical(idx.append(idx), categories=idx)
296
+ exp = """[2011-01, 2011-02, 2011-03, 2011-04, 2011-05, 2011-01, 2011-02, 2011-03, 2011-04, 2011-05]
297
+ Categories (5, period[M]): [2011-01, 2011-02, 2011-03, 2011-04, 2011-05]""" # noqa: E501
298
+
299
+ assert repr(c) == exp
300
+
301
+ def test_categorical_repr_period_ordered(self):
302
+ idx = period_range("2011-01-01 09:00", freq="h", periods=5)
303
+ c = Categorical(idx, ordered=True)
304
+ exp = """[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00]
305
+ Categories (5, period[h]): [2011-01-01 09:00 < 2011-01-01 10:00 < 2011-01-01 11:00 < 2011-01-01 12:00 <
306
+ 2011-01-01 13:00]""" # noqa: E501
307
+
308
+ assert repr(c) == exp
309
+
310
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
311
+ exp = """[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00, 2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00]
312
+ Categories (5, period[h]): [2011-01-01 09:00 < 2011-01-01 10:00 < 2011-01-01 11:00 < 2011-01-01 12:00 <
313
+ 2011-01-01 13:00]""" # noqa: E501
314
+
315
+ assert repr(c) == exp
316
+
317
+ idx = period_range("2011-01", freq="M", periods=5)
318
+ c = Categorical(idx, ordered=True)
319
+ exp = """[2011-01, 2011-02, 2011-03, 2011-04, 2011-05]
320
+ Categories (5, period[M]): [2011-01 < 2011-02 < 2011-03 < 2011-04 < 2011-05]"""
321
+
322
+ assert repr(c) == exp
323
+
324
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
325
+ exp = """[2011-01, 2011-02, 2011-03, 2011-04, 2011-05, 2011-01, 2011-02, 2011-03, 2011-04, 2011-05]
326
+ Categories (5, period[M]): [2011-01 < 2011-02 < 2011-03 < 2011-04 < 2011-05]""" # noqa: E501
327
+
328
+ assert repr(c) == exp
329
+
330
+ def test_categorical_repr_timedelta(self):
331
+ idx = timedelta_range("1 days", periods=5)
332
+ c = Categorical(idx)
333
+ exp = """[1 days, 2 days, 3 days, 4 days, 5 days]
334
+ Categories (5, timedelta64[us]): [1 days, 2 days, 3 days, 4 days, 5 days]"""
335
+
336
+ assert repr(c) == exp
337
+
338
+ c = Categorical(idx.append(idx), categories=idx)
339
+ exp = """[1 days, 2 days, 3 days, 4 days, 5 days, 1 days, 2 days, 3 days, 4 days, 5 days]
340
+ Categories (5, timedelta64[us]): [1 days, 2 days, 3 days, 4 days, 5 days]""" # noqa: E501
341
+
342
+ assert repr(c) == exp
343
+
344
+ idx = timedelta_range("1 hours", periods=20)
345
+ c = Categorical(idx)
346
+ exp = """[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, 4 days 01:00:00, ..., 15 days 01:00:00, 16 days 01:00:00, 17 days 01:00:00, 18 days 01:00:00, 19 days 01:00:00]
347
+ Length: 20
348
+ Categories (20, timedelta64[us]): [0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00,
349
+ 3 days 01:00:00, ..., 16 days 01:00:00, 17 days 01:00:00,
350
+ 18 days 01:00:00, 19 days 01:00:00]""" # noqa: E501
351
+
352
+ assert repr(c) == exp
353
+
354
+ c = Categorical(idx.append(idx), categories=idx)
355
+ exp = """[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, 4 days 01:00:00, ..., 15 days 01:00:00, 16 days 01:00:00, 17 days 01:00:00, 18 days 01:00:00, 19 days 01:00:00]
356
+ Length: 40
357
+ Categories (20, timedelta64[us]): [0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00,
358
+ 3 days 01:00:00, ..., 16 days 01:00:00, 17 days 01:00:00,
359
+ 18 days 01:00:00, 19 days 01:00:00]""" # noqa: E501
360
+
361
+ assert repr(c) == exp
362
+
363
+ def test_categorical_repr_timedelta_ordered(self):
364
+ idx = timedelta_range("1 days", periods=5)
365
+ c = Categorical(idx, ordered=True)
366
+ exp = """[1 days, 2 days, 3 days, 4 days, 5 days]
367
+ Categories (5, timedelta64[us]): [1 days < 2 days < 3 days < 4 days < 5 days]"""
368
+
369
+ assert repr(c) == exp
370
+
371
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
372
+ exp = """[1 days, 2 days, 3 days, 4 days, 5 days, 1 days, 2 days, 3 days, 4 days, 5 days]
373
+ Categories (5, timedelta64[us]): [1 days < 2 days < 3 days < 4 days < 5 days]""" # noqa: E501
374
+
375
+ assert repr(c) == exp
376
+
377
+ idx = timedelta_range("1 hours", periods=20)
378
+ c = Categorical(idx, ordered=True)
379
+ exp = """[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, 4 days 01:00:00, ..., 15 days 01:00:00, 16 days 01:00:00, 17 days 01:00:00, 18 days 01:00:00, 19 days 01:00:00]
380
+ Length: 20
381
+ Categories (20, timedelta64[us]): [0 days 01:00:00 < 1 days 01:00:00 < 2 days 01:00:00 <
382
+ 3 days 01:00:00 ... 16 days 01:00:00 < 17 days 01:00:00 <
383
+ 18 days 01:00:00 < 19 days 01:00:00]""" # noqa: E501
384
+
385
+ assert repr(c) == exp
386
+
387
+ c = Categorical(idx.append(idx), categories=idx, ordered=True)
388
+ exp = """[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, 4 days 01:00:00, ..., 15 days 01:00:00, 16 days 01:00:00, 17 days 01:00:00, 18 days 01:00:00, 19 days 01:00:00]
389
+ Length: 40
390
+ Categories (20, timedelta64[us]): [0 days 01:00:00 < 1 days 01:00:00 < 2 days 01:00:00 <
391
+ 3 days 01:00:00 ... 16 days 01:00:00 < 17 days 01:00:00 <
392
+ 18 days 01:00:00 < 19 days 01:00:00]""" # noqa: E501
393
+
394
+ assert repr(c) == exp
395
+
396
+ def test_categorical_index_repr(self):
397
+ idx = CategoricalIndex(Categorical([1, 2, 3]))
398
+ exp = """CategoricalIndex([1, 2, 3], categories=[1, 2, 3], ordered=False, dtype='category')""" # noqa: E501
399
+ assert repr(idx) == exp
400
+
401
+ i = CategoricalIndex(Categorical(np.arange(10, dtype=np.int64)))
402
+ exp = """CategoricalIndex([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], categories=[0, 1, 2, 3, ..., 6, 7, 8, 9], ordered=False, dtype='category')""" # noqa: E501
403
+ assert repr(i) == exp
404
+
405
+ def test_categorical_index_repr_ordered(self):
406
+ i = CategoricalIndex(Categorical([1, 2, 3], ordered=True))
407
+ exp = """CategoricalIndex([1, 2, 3], categories=[1, 2, 3], ordered=True, dtype='category')""" # noqa: E501
408
+ assert repr(i) == exp
409
+
410
+ i = CategoricalIndex(Categorical(np.arange(10, dtype=np.int64), ordered=True))
411
+ exp = """CategoricalIndex([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], categories=[0, 1, 2, 3, ..., 6, 7, 8, 9], ordered=True, dtype='category')""" # noqa: E501
412
+ assert repr(i) == exp
413
+
414
+ def test_categorical_index_repr_datetime(self):
415
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5)
416
+ i = CategoricalIndex(Categorical(idx))
417
+ exp = """CategoricalIndex(['2011-01-01 09:00:00', '2011-01-01 10:00:00',
418
+ '2011-01-01 11:00:00', '2011-01-01 12:00:00',
419
+ '2011-01-01 13:00:00'],
420
+ categories=[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, 2011-01-01 13:00:00], ordered=False, dtype='category')""" # noqa: E501
421
+
422
+ assert repr(i) == exp
423
+
424
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5, tz="US/Eastern")
425
+ i = CategoricalIndex(Categorical(idx))
426
+ exp = """CategoricalIndex(['2011-01-01 09:00:00-05:00', '2011-01-01 10:00:00-05:00',
427
+ '2011-01-01 11:00:00-05:00', '2011-01-01 12:00:00-05:00',
428
+ '2011-01-01 13:00:00-05:00'],
429
+ categories=[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00], ordered=False, dtype='category')""" # noqa: E501
430
+
431
+ assert repr(i) == exp
432
+
433
+ def test_categorical_index_repr_datetime_ordered(self):
434
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5)
435
+ i = CategoricalIndex(Categorical(idx, ordered=True))
436
+ exp = """CategoricalIndex(['2011-01-01 09:00:00', '2011-01-01 10:00:00',
437
+ '2011-01-01 11:00:00', '2011-01-01 12:00:00',
438
+ '2011-01-01 13:00:00'],
439
+ categories=[2011-01-01 09:00:00, 2011-01-01 10:00:00, 2011-01-01 11:00:00, 2011-01-01 12:00:00, 2011-01-01 13:00:00], ordered=True, dtype='category')""" # noqa: E501
440
+
441
+ assert repr(i) == exp
442
+
443
+ idx = date_range("2011-01-01 09:00", freq="h", periods=5, tz="US/Eastern")
444
+ i = CategoricalIndex(Categorical(idx, ordered=True))
445
+ exp = """CategoricalIndex(['2011-01-01 09:00:00-05:00', '2011-01-01 10:00:00-05:00',
446
+ '2011-01-01 11:00:00-05:00', '2011-01-01 12:00:00-05:00',
447
+ '2011-01-01 13:00:00-05:00'],
448
+ categories=[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00], ordered=True, dtype='category')""" # noqa: E501
449
+
450
+ assert repr(i) == exp
451
+
452
+ i = CategoricalIndex(Categorical(idx.append(idx), ordered=True))
453
+ exp = """CategoricalIndex(['2011-01-01 09:00:00-05:00', '2011-01-01 10:00:00-05:00',
454
+ '2011-01-01 11:00:00-05:00', '2011-01-01 12:00:00-05:00',
455
+ '2011-01-01 13:00:00-05:00', '2011-01-01 09:00:00-05:00',
456
+ '2011-01-01 10:00:00-05:00', '2011-01-01 11:00:00-05:00',
457
+ '2011-01-01 12:00:00-05:00', '2011-01-01 13:00:00-05:00'],
458
+ categories=[2011-01-01 09:00:00-05:00, 2011-01-01 10:00:00-05:00, 2011-01-01 11:00:00-05:00, 2011-01-01 12:00:00-05:00, 2011-01-01 13:00:00-05:00], ordered=True, dtype='category')""" # noqa: E501
459
+
460
+ assert repr(i) == exp
461
+
462
+ def test_categorical_index_repr_period(self):
463
+ # test all length
464
+ idx = period_range("2011-01-01 09:00", freq="h", periods=1)
465
+ i = CategoricalIndex(Categorical(idx))
466
+ exp = """CategoricalIndex(['2011-01-01 09:00'], categories=[2011-01-01 09:00], ordered=False, dtype='category')""" # noqa: E501
467
+ assert repr(i) == exp
468
+
469
+ idx = period_range("2011-01-01 09:00", freq="h", periods=2)
470
+ i = CategoricalIndex(Categorical(idx))
471
+ exp = """CategoricalIndex(['2011-01-01 09:00', '2011-01-01 10:00'], categories=[2011-01-01 09:00, 2011-01-01 10:00], ordered=False, dtype='category')""" # noqa: E501
472
+ assert repr(i) == exp
473
+
474
+ idx = period_range("2011-01-01 09:00", freq="h", periods=3)
475
+ i = CategoricalIndex(Categorical(idx))
476
+ exp = """CategoricalIndex(['2011-01-01 09:00', '2011-01-01 10:00', '2011-01-01 11:00'], categories=[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00], ordered=False, dtype='category')""" # noqa: E501
477
+ assert repr(i) == exp
478
+
479
+ idx = period_range("2011-01-01 09:00", freq="h", periods=5)
480
+ i = CategoricalIndex(Categorical(idx))
481
+ exp = """CategoricalIndex(['2011-01-01 09:00', '2011-01-01 10:00', '2011-01-01 11:00',
482
+ '2011-01-01 12:00', '2011-01-01 13:00'],
483
+ categories=[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00], ordered=False, dtype='category')""" # noqa: E501
484
+
485
+ assert repr(i) == exp
486
+
487
+ i = CategoricalIndex(Categorical(idx.append(idx)))
488
+ exp = """CategoricalIndex(['2011-01-01 09:00', '2011-01-01 10:00', '2011-01-01 11:00',
489
+ '2011-01-01 12:00', '2011-01-01 13:00', '2011-01-01 09:00',
490
+ '2011-01-01 10:00', '2011-01-01 11:00', '2011-01-01 12:00',
491
+ '2011-01-01 13:00'],
492
+ categories=[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00], ordered=False, dtype='category')""" # noqa: E501
493
+
494
+ assert repr(i) == exp
495
+
496
+ idx = period_range("2011-01", freq="M", periods=5)
497
+ i = CategoricalIndex(Categorical(idx))
498
+ exp = """CategoricalIndex(['2011-01', '2011-02', '2011-03', '2011-04', '2011-05'], categories=[2011-01, 2011-02, 2011-03, 2011-04, 2011-05], ordered=False, dtype='category')""" # noqa: E501
499
+ assert repr(i) == exp
500
+
501
+ def test_categorical_index_repr_period_ordered(self):
502
+ idx = period_range("2011-01-01 09:00", freq="h", periods=5)
503
+ i = CategoricalIndex(Categorical(idx, ordered=True))
504
+ exp = """CategoricalIndex(['2011-01-01 09:00', '2011-01-01 10:00', '2011-01-01 11:00',
505
+ '2011-01-01 12:00', '2011-01-01 13:00'],
506
+ categories=[2011-01-01 09:00, 2011-01-01 10:00, 2011-01-01 11:00, 2011-01-01 12:00, 2011-01-01 13:00], ordered=True, dtype='category')""" # noqa: E501
507
+
508
+ assert repr(i) == exp
509
+
510
+ idx = period_range("2011-01", freq="M", periods=5)
511
+ i = CategoricalIndex(Categorical(idx, ordered=True))
512
+ exp = """CategoricalIndex(['2011-01', '2011-02', '2011-03', '2011-04', '2011-05'], categories=[2011-01, 2011-02, 2011-03, 2011-04, 2011-05], ordered=True, dtype='category')""" # noqa: E501
513
+ assert repr(i) == exp
514
+
515
+ def test_categorical_index_repr_timedelta(self):
516
+ idx = timedelta_range("1 days", periods=5)
517
+ i = CategoricalIndex(Categorical(idx))
518
+ exp = """CategoricalIndex(['1 days', '2 days', '3 days', '4 days', '5 days'], categories=[1 days, 2 days, 3 days, 4 days, 5 days], ordered=False, dtype='category')""" # noqa: E501
519
+ assert repr(i) == exp
520
+
521
+ idx = timedelta_range("1 hours", periods=10)
522
+ i = CategoricalIndex(Categorical(idx))
523
+ exp = """CategoricalIndex(['0 days 01:00:00', '1 days 01:00:00', '2 days 01:00:00',
524
+ '3 days 01:00:00', '4 days 01:00:00', '5 days 01:00:00',
525
+ '6 days 01:00:00', '7 days 01:00:00', '8 days 01:00:00',
526
+ '9 days 01:00:00'],
527
+ categories=[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, ..., 6 days 01:00:00, 7 days 01:00:00, 8 days 01:00:00, 9 days 01:00:00], ordered=False, dtype='category')""" # noqa: E501
528
+
529
+ assert repr(i) == exp
530
+
531
+ def test_categorical_index_repr_timedelta_ordered(self):
532
+ idx = timedelta_range("1 days", periods=5)
533
+ i = CategoricalIndex(Categorical(idx, ordered=True))
534
+ exp = """CategoricalIndex(['1 days', '2 days', '3 days', '4 days', '5 days'], categories=[1 days, 2 days, 3 days, 4 days, 5 days], ordered=True, dtype='category')""" # noqa: E501
535
+ assert repr(i) == exp
536
+
537
+ idx = timedelta_range("1 hours", periods=10)
538
+ i = CategoricalIndex(Categorical(idx, ordered=True))
539
+ exp = """CategoricalIndex(['0 days 01:00:00', '1 days 01:00:00', '2 days 01:00:00',
540
+ '3 days 01:00:00', '4 days 01:00:00', '5 days 01:00:00',
541
+ '6 days 01:00:00', '7 days 01:00:00', '8 days 01:00:00',
542
+ '9 days 01:00:00'],
543
+ categories=[0 days 01:00:00, 1 days 01:00:00, 2 days 01:00:00, 3 days 01:00:00, ..., 6 days 01:00:00, 7 days 01:00:00, 8 days 01:00:00, 9 days 01:00:00], ordered=True, dtype='category')""" # noqa: E501
544
+
545
+ assert repr(i) == exp
546
+
547
+ def test_categorical_str_repr(self):
548
+ # GH 33676
549
+ result = repr(Categorical([1, "2", 3, 4]))
550
+ expected = "[1, '2', 3, 4]\nCategories (4, object): [1, 3, 4, '2']"
551
+ assert result == expected
552
+
553
+ def test_categorical_with_string_dtype(self, string_dtype_no_object):
554
+ # GH 63045 - ensure categories are quoted for string dtypes
555
+ s = Series(
556
+ ["apple", "banana", "cherry", "cherry"], dtype=string_dtype_no_object
557
+ )
558
+ result = repr(Categorical(s))
559
+ expected = f"['apple', 'banana', 'cherry', 'cherry']\nCategories (3, {string_dtype_no_object!s}): ['apple', 'banana', 'cherry']" # noqa: E501
560
+
561
+ assert result == expected
.venv/lib/python3.13/site-packages/pandas/tests/arrays/categorical/test_sorting.py ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from pandas import (
5
+ Categorical,
6
+ Index,
7
+ )
8
+ import pandas._testing as tm
9
+
10
+
11
+ class TestCategoricalSort:
12
+ def test_argsort(self):
13
+ c = Categorical([5, 3, 1, 4, 2], ordered=True)
14
+
15
+ expected = np.array([2, 4, 1, 3, 0])
16
+ tm.assert_numpy_array_equal(
17
+ c.argsort(ascending=True), expected, check_dtype=False
18
+ )
19
+
20
+ expected = expected[::-1]
21
+ tm.assert_numpy_array_equal(
22
+ c.argsort(ascending=False), expected, check_dtype=False
23
+ )
24
+
25
+ def test_numpy_argsort(self):
26
+ c = Categorical([5, 3, 1, 4, 2], ordered=True)
27
+
28
+ expected = np.array([2, 4, 1, 3, 0])
29
+ tm.assert_numpy_array_equal(np.argsort(c), expected, check_dtype=False)
30
+
31
+ tm.assert_numpy_array_equal(
32
+ np.argsort(c, kind="mergesort"), expected, check_dtype=False
33
+ )
34
+
35
+ msg = "the 'axis' parameter is not supported"
36
+ with pytest.raises(ValueError, match=msg):
37
+ np.argsort(c, axis=0)
38
+
39
+ msg = "the 'order' parameter is not supported"
40
+ with pytest.raises(ValueError, match=msg):
41
+ np.argsort(c, order="C")
42
+
43
+ def test_sort_values(self):
44
+ # unordered cats are sortable
45
+ cat = Categorical(["a", "b", "b", "a"], ordered=False)
46
+ cat.sort_values()
47
+
48
+ cat = Categorical(["a", "c", "b", "d"], ordered=True)
49
+
50
+ # sort_values
51
+ res = cat.sort_values()
52
+ exp = np.array(["a", "b", "c", "d"], dtype=object)
53
+ tm.assert_numpy_array_equal(res.__array__(), exp)
54
+ tm.assert_index_equal(res.categories, cat.categories)
55
+
56
+ cat = Categorical(
57
+ ["a", "c", "b", "d"], categories=["a", "b", "c", "d"], ordered=True
58
+ )
59
+ res = cat.sort_values()
60
+ exp = np.array(["a", "b", "c", "d"], dtype=object)
61
+ tm.assert_numpy_array_equal(res.__array__(), exp)
62
+ tm.assert_index_equal(res.categories, cat.categories)
63
+
64
+ res = cat.sort_values(ascending=False)
65
+ exp = np.array(["d", "c", "b", "a"], dtype=object)
66
+ tm.assert_numpy_array_equal(res.__array__(), exp)
67
+ tm.assert_index_equal(res.categories, cat.categories)
68
+
69
+ # sort (inplace order)
70
+ cat1 = cat.copy()
71
+ orig_codes = cat1._codes
72
+ cat1.sort_values(inplace=True)
73
+ assert cat1._codes is orig_codes
74
+ exp = np.array(["a", "b", "c", "d"], dtype=object)
75
+ tm.assert_numpy_array_equal(cat1.__array__(), exp)
76
+ tm.assert_index_equal(res.categories, cat.categories)
77
+
78
+ # reverse
79
+ cat = Categorical(["a", "c", "c", "b", "d"], ordered=True)
80
+ res = cat.sort_values(ascending=False)
81
+ exp_val = np.array(["d", "c", "c", "b", "a"], dtype=object)
82
+ exp_categories = Index(["a", "b", "c", "d"])
83
+ tm.assert_numpy_array_equal(res.__array__(), exp_val)
84
+ tm.assert_index_equal(res.categories, exp_categories)
85
+
86
+ def test_sort_values_na_position(self):
87
+ # see gh-12882
88
+ cat = Categorical([5, 2, np.nan, 2, np.nan], ordered=True)
89
+ exp_categories = Index([2, 5])
90
+
91
+ exp = np.array([2.0, 2.0, 5.0, np.nan, np.nan])
92
+ res = cat.sort_values() # default arguments
93
+ tm.assert_numpy_array_equal(res.__array__(), exp)
94
+ tm.assert_index_equal(res.categories, exp_categories)
95
+
96
+ exp = np.array([np.nan, np.nan, 2.0, 2.0, 5.0])
97
+ res = cat.sort_values(ascending=True, na_position="first")
98
+ tm.assert_numpy_array_equal(res.__array__(), exp)
99
+ tm.assert_index_equal(res.categories, exp_categories)
100
+
101
+ exp = np.array([np.nan, np.nan, 5.0, 2.0, 2.0])
102
+ res = cat.sort_values(ascending=False, na_position="first")
103
+ tm.assert_numpy_array_equal(res.__array__(), exp)
104
+ tm.assert_index_equal(res.categories, exp_categories)
105
+
106
+ exp = np.array([2.0, 2.0, 5.0, np.nan, np.nan])
107
+ res = cat.sort_values(ascending=True, na_position="last")
108
+ tm.assert_numpy_array_equal(res.__array__(), exp)
109
+ tm.assert_index_equal(res.categories, exp_categories)
110
+
111
+ exp = np.array([5.0, 2.0, 2.0, np.nan, np.nan])
112
+ res = cat.sort_values(ascending=False, na_position="last")
113
+ tm.assert_numpy_array_equal(res.__array__(), exp)
114
+ tm.assert_index_equal(res.categories, exp_categories)
115
+
116
+ cat = Categorical(["a", "c", "b", "d", np.nan], ordered=True)
117
+ res = cat.sort_values(ascending=False, na_position="last")
118
+ exp_val = np.array(["d", "c", "b", "a", np.nan], dtype=object)
119
+ exp_categories = Index(["a", "b", "c", "d"])
120
+ tm.assert_numpy_array_equal(res.__array__(), exp_val)
121
+ tm.assert_index_equal(res.categories, exp_categories)
122
+
123
+ cat = Categorical(["a", "c", "b", "d", np.nan], ordered=True)
124
+ res = cat.sort_values(ascending=False, na_position="first")
125
+ exp_val = np.array([np.nan, "d", "c", "b", "a"], dtype=object)
126
+ exp_categories = Index(["a", "b", "c", "d"])
127
+ tm.assert_numpy_array_equal(res.__array__(), exp_val)
128
+ tm.assert_index_equal(res.categories, exp_categories)