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