1"""
2datetimelke_accumulations.py is for accumulations of datetimelike extension arrays
3"""
4
5from __future__ import annotations
6
7from typing import TYPE_CHECKING
8
9import numpy as np
10
11from pandas._libs import iNaT
12
13from pandas.core.dtypes.missing import isna
14
15if TYPE_CHECKING:
16 from collections.abc import Callable
17
18
19def _cum_func(
20 func: Callable,
21 values: np.ndarray,
22 *,
23 skipna: bool = True,
24) -> np.ndarray:
25 """
26 Accumulations for 1D datetimelike arrays.
27
28 Parameters
29 ----------
30 func : np.cumsum, np.maximum.accumulate, np.minimum.accumulate
31 values : np.ndarray
32 Numpy array with the values (can be of any dtype that support the
33 operation). Values is changed is modified inplace.
34 skipna : bool, default True
35 Whether to skip NA.
36 """
37 try:
38 fill_value = {
39 np.maximum.accumulate: np.iinfo(np.int64).min,
40 np.cumsum: 0,
41 np.minimum.accumulate: np.iinfo(np.int64).max,
42 }[func]
43 except KeyError as err:
44 raise ValueError(
45 f"No accumulation for {func} implemented on BaseMaskedArray"
46 ) from err
47
48 mask = isna(values)
49 y = values.view("i8")
50 y[mask] = fill_value
51
52 if not skipna:
53 mask = np.maximum.accumulate(mask)
54
55 # GH 57956
56 result = func(y, axis=0)
57 result[mask] = iNaT
58
59 if values.dtype.kind in "mM":
60 return result.view(values.dtype.base)
61 return result
62
63
64def cumsum(values: np.ndarray, *, skipna: bool = True) -> np.ndarray:
65 return _cum_func(np.cumsum, values, skipna=skipna)
66
67
68def cummin(values: np.ndarray, *, skipna: bool = True) -> np.ndarray:
69 return _cum_func(np.minimum.accumulate, values, skipna=skipna)
70
71
72def cummax(values: np.ndarray, *, skipna: bool = True) -> np.ndarray:
73 return _cum_func(np.maximum.accumulate, values, skipna=skipna)