1from __future__ import annotations
2
3from typing import (
4 TYPE_CHECKING,
5 Any,
6 Literal,
7 Self,
8 cast,
9)
10
11import numpy as np
12
13from pandas._libs import lib
14from pandas._libs.tslibs import is_supported_dtype
15from pandas.compat.numpy import function as nv
16from pandas.util._decorators import set_module
17
18from pandas.core.dtypes.astype import (
19 astype_array,
20 astype_is_view,
21)
22from pandas.core.dtypes.cast import (
23 construct_1d_object_array_from_listlike,
24 maybe_downcast_to_dtype,
25)
26from pandas.core.dtypes.common import pandas_dtype
27from pandas.core.dtypes.dtypes import NumpyEADtype
28from pandas.core.dtypes.missing import isna
29
30from pandas.core import (
31 arraylike,
32 missing,
33 nanops,
34 ops,
35)
36from pandas.core.arraylike import OpsMixin
37from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
38from pandas.core.construction import ensure_wrapped_if_datetimelike
39from pandas.core.strings.object_array import ObjectStringArrayMixin
40
41if TYPE_CHECKING:
42 from collections.abc import Callable
43
44 from pandas._typing import (
45 ArrayLike,
46 AxisInt,
47 Dtype,
48 FillnaOptions,
49 InterpolateOptions,
50 NpDtype,
51 Scalar,
52 TakeIndexer,
53 npt,
54 )
55
56 from pandas import Index
57 from pandas.arrays import StringArray
58
59
60@set_module("pandas.arrays")
61class NumpyExtensionArray(
62 OpsMixin,
63 NDArrayBackedExtensionArray,
64 ObjectStringArrayMixin,
65):
66 """
67 A pandas ExtensionArray for NumPy data.
68
69 This is mostly for internal compatibility, and is not especially
70 useful on its own.
71
72 Parameters
73 ----------
74 values : ndarray
75 The NumPy ndarray to wrap. Must be 1-dimensional.
76 copy : bool, default False
77 Whether to copy `values`.
78
79 Attributes
80 ----------
81 None
82
83 Methods
84 -------
85 None
86
87 See Also
88 --------
89 array : Create an array.
90 Series.to_numpy : Convert a Series to a NumPy array.
91
92 Examples
93 --------
94 >>> pd.arrays.NumpyExtensionArray(np.array([0, 1, 2, 3]))
95 <NumpyExtensionArray>
96 [0, 1, 2, 3]
97 Length: 4, dtype: int64
98 """
99
100 # If you're wondering why pd.Series(cls) doesn't put the array in an
101 # ExtensionBlock, search for `ABCNumpyExtensionArray`. We check for
102 # that _typ to ensure that users don't unnecessarily use EAs inside
103 # pandas internals, which turns off things like block consolidation.
104 _typ = "npy_extension"
105 __array_priority__ = 1000
106 _ndarray: np.ndarray
107 _dtype: NumpyEADtype
108 _internal_fill_value = np.nan
109
110 # ------------------------------------------------------------------------
111 # Constructors
112
113 def __init__(
114 self, values: np.ndarray | NumpyExtensionArray, copy: bool = False
115 ) -> None:
116 if isinstance(values, type(self)):
117 values = values._ndarray
118 if not isinstance(values, np.ndarray):
119 raise ValueError(
120 f"'values' must be a NumPy array, not {type(values).__name__}"
121 )
122
123 if values.ndim == 0:
124 # Technically we support 2, but do not advertise that fact.
125 raise ValueError("NumpyExtensionArray must be 1-dimensional.")
126
127 if copy:
128 values = values.copy()
129
130 dtype = NumpyEADtype(values.dtype)
131 super().__init__(values, dtype)
132
133 @classmethod
134 def _from_sequence(
135 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
136 ) -> NumpyExtensionArray:
137 if isinstance(dtype, NumpyEADtype):
138 dtype = dtype._dtype
139
140 # error: Argument "dtype" to "asarray" has incompatible type
141 # "Union[ExtensionDtype, str, dtype[Any], dtype[floating[_64Bit]], Type[object],
142 # None]"; expected "Union[dtype[Any], None, type, _SupportsDType, str,
143 # Union[Tuple[Any, int], Tuple[Any, Union[int, Sequence[int]]], List[Any],
144 # _DTypeDict, Tuple[Any, Any]]]"
145 result = np.asarray(scalars, dtype=dtype) # type: ignore[arg-type]
146 if (
147 result.ndim > 1
148 and not hasattr(scalars, "dtype")
149 and (dtype is None or dtype == object)
150 ):
151 # e.g. list-of-tuples
152 result = construct_1d_object_array_from_listlike(scalars)
153
154 if copy and result is scalars:
155 result = result.copy()
156 return cls(result)
157
158 def _cast_pointwise_result(self, values) -> ArrayLike:
159 result = super()._cast_pointwise_result(values)
160 lkind = self.dtype.kind
161 rkind = result.dtype.kind
162 if (
163 (lkind in "iu" and rkind in "iu")
164 or (lkind == "f" and rkind == "f")
165 or (lkind == rkind == "c")
166 ):
167 result = maybe_downcast_to_dtype(result, self.dtype.numpy_dtype)
168 elif rkind == "M":
169 # Ensure potential subsequent .astype(object) doesn't incorrectly
170 # convert Timestamps to ints
171 from pandas import array as pd_array
172
173 result = pd_array(result, copy=False)
174 return result
175
176 # ------------------------------------------------------------------------
177 # Data
178
179 @property
180 def dtype(self) -> NumpyEADtype:
181 return self._dtype
182
183 # ------------------------------------------------------------------------
184 # NumPy Array Interface
185
186 def __array__(
187 self, dtype: np.dtype | None = None, copy: bool | None = None
188 ) -> np.ndarray:
189 if copy is not None:
190 # Note: branch avoids `copy=None` for NumPy 1.x support
191 result = np.array(self._ndarray, dtype=dtype, copy=copy)
192 else:
193 result = np.asarray(self._ndarray, dtype=dtype)
194
195 if (
196 self._readonly
197 and not copy
198 and (dtype is None or astype_is_view(self.dtype, dtype))
199 ):
200 result = result.view()
201 result.flags.writeable = False
202
203 return result
204
205 def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
206 # Lightly modified version of
207 # https://numpy.org/doc/stable/reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html
208 # The primary modification is not boxing scalar return values
209 # in NumpyExtensionArray, since pandas' ExtensionArrays are 1-d.
210 out = kwargs.get("out", ())
211
212 result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
213 self, ufunc, method, *inputs, **kwargs
214 )
215 if result is not NotImplemented:
216 return result
217
218 if "out" in kwargs:
219 # e.g. test_ufunc_unary
220 return arraylike.dispatch_ufunc_with_out(
221 self, ufunc, method, *inputs, **kwargs
222 )
223
224 if method == "reduce":
225 result = arraylike.dispatch_reduction_ufunc(
226 self, ufunc, method, *inputs, **kwargs
227 )
228 if result is not NotImplemented:
229 # e.g. tests.series.test_ufunc.TestNumpyReductions
230 return result
231
232 # Defer to the implementation of the ufunc on unwrapped values.
233 inputs = tuple(
234 x._ndarray if isinstance(x, NumpyExtensionArray) else x for x in inputs
235 )
236 if out:
237 kwargs["out"] = tuple(
238 x._ndarray if isinstance(x, NumpyExtensionArray) else x for x in out
239 )
240 result = getattr(ufunc, method)(*inputs, **kwargs)
241
242 if ufunc.nout > 1:
243 # multiple return values; re-box array-like results
244 return tuple(type(self)(x) for x in result)
245 elif method == "at":
246 # no return value
247 return None
248 elif method == "reduce":
249 if isinstance(result, np.ndarray):
250 # e.g. test_np_reduce_2d
251 return type(self)(result)
252
253 # e.g. test_np_max_nested_tuples
254 return result
255 else:
256 if self.dtype.type is str: # type: ignore[comparison-overlap]
257 # StringDtype
258 self = cast("StringArray", self)
259 try:
260 # specify dtype to preserve storage/na_value
261 return type(self)(result, dtype=self.dtype)
262 except ValueError:
263 # if validation of input fails (no strings)
264 # -> fallback to returning raw numpy array
265 return result
266 # one return value; re-box array-like results
267 return type(self)(result)
268
269 # ------------------------------------------------------------------------
270 # Pandas ExtensionArray Interface
271
272 def astype(self, dtype, copy: bool = True):
273 dtype = pandas_dtype(dtype)
274
275 if dtype == self.dtype:
276 if copy:
277 return self.copy()
278 return self
279
280 result = astype_array(self._ndarray, dtype=dtype, copy=copy)
281 return result
282
283 def isna(self) -> np.ndarray:
284 return isna(self._ndarray)
285
286 def _validate_scalar(self, fill_value):
287 if fill_value is None:
288 # Primarily for subclasses
289 fill_value = self.dtype.na_value
290 return fill_value
291
292 def _values_for_factorize(self) -> tuple[np.ndarray, float | None]:
293 if self.dtype.kind in "iub":
294 fv = None
295 else:
296 fv = np.nan
297 return self._ndarray, fv
298
299 # Base EA class (and all other EA classes) don't have limit_area keyword
300 # This can be removed here as well when the interpolate ffill/bfill method
301 # deprecation is enforced
302 def _pad_or_backfill(
303 self,
304 *,
305 method: FillnaOptions,
306 limit: int | None = None,
307 limit_area: Literal["inside", "outside"] | None = None,
308 copy: bool = True,
309 ) -> Self:
310 """
311 ffill or bfill along axis=0.
312 """
313 if copy:
314 out_data = self._ndarray.copy()
315 else:
316 out_data = self._ndarray
317
318 meth = missing.clean_fill_method(method)
319 missing.pad_or_backfill_inplace(
320 out_data.T,
321 method=meth,
322 axis=0,
323 limit=limit,
324 limit_area=limit_area,
325 )
326
327 if not copy:
328 return self
329 return type(self)._simple_new(out_data, dtype=self.dtype)
330
331 def interpolate(
332 self,
333 *,
334 method: InterpolateOptions,
335 axis: int,
336 index: Index,
337 limit,
338 limit_direction,
339 limit_area,
340 copy: bool,
341 **kwargs,
342 ) -> Self:
343 """
344 See NDFrame.interpolate.__doc__.
345 """
346 # NB: we return type(self) even if copy=False
347 if not self.dtype._is_numeric:
348 raise TypeError(f"Cannot interpolate with {self.dtype} dtype")
349
350 if not copy:
351 out_data = self._ndarray
352 else:
353 out_data = self._ndarray.copy()
354
355 # TODO: assert we have floating dtype?
356 missing.interpolate_2d_inplace(
357 out_data,
358 method=method,
359 axis=axis,
360 index=index,
361 limit=limit,
362 limit_direction=limit_direction,
363 limit_area=limit_area,
364 **kwargs,
365 )
366 if not copy:
367 return self
368 return type(self)._simple_new(out_data, dtype=self.dtype)
369
370 def take(
371 self,
372 indices: TakeIndexer,
373 *,
374 allow_fill: bool = False,
375 fill_value: Any = None,
376 axis: AxisInt = 0,
377 ) -> Self:
378 """
379 Take entries from this array at each index in a list of indices,
380 producing an array containing only those entries.
381 """
382 result = super().take(
383 indices, allow_fill=allow_fill, fill_value=fill_value, axis=axis
384 )
385 # See GH#62448.
386 if self.dtype.kind in "iub":
387 return type(self)(result._ndarray, copy=False)
388
389 return result
390
391 # ------------------------------------------------------------------------
392 # Reductions
393
394 def any(
395 self,
396 *,
397 axis: AxisInt | None = None,
398 out=None,
399 keepdims: bool = False,
400 skipna: bool = True,
401 ):
402 nv.validate_any((), {"out": out, "keepdims": keepdims})
403 result = nanops.nanany(self._ndarray, axis=axis, skipna=skipna)
404 return self._wrap_reduction_result(axis, result)
405
406 def all(
407 self,
408 *,
409 axis: AxisInt | None = None,
410 out=None,
411 keepdims: bool = False,
412 skipna: bool = True,
413 ):
414 nv.validate_all((), {"out": out, "keepdims": keepdims})
415 result = nanops.nanall(self._ndarray, axis=axis, skipna=skipna)
416 return self._wrap_reduction_result(axis, result)
417
418 def min(
419 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
420 ) -> Scalar:
421 nv.validate_min((), kwargs)
422 result = nanops.nanmin(
423 values=self._ndarray, axis=axis, mask=self.isna(), skipna=skipna
424 )
425 return self._wrap_reduction_result(axis, result)
426
427 def max(
428 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
429 ) -> Scalar:
430 nv.validate_max((), kwargs)
431 result = nanops.nanmax(
432 values=self._ndarray, axis=axis, mask=self.isna(), skipna=skipna
433 )
434 return self._wrap_reduction_result(axis, result)
435
436 def sum(
437 self,
438 *,
439 axis: AxisInt | None = None,
440 skipna: bool = True,
441 min_count: int = 0,
442 **kwargs,
443 ) -> Scalar:
444 nv.validate_sum((), kwargs)
445 result = nanops.nansum(
446 self._ndarray, axis=axis, skipna=skipna, min_count=min_count
447 )
448 return self._wrap_reduction_result(axis, result)
449
450 def prod(
451 self,
452 *,
453 axis: AxisInt | None = None,
454 skipna: bool = True,
455 min_count: int = 0,
456 **kwargs,
457 ) -> Scalar:
458 nv.validate_prod((), kwargs)
459 result = nanops.nanprod(
460 self._ndarray, axis=axis, skipna=skipna, min_count=min_count
461 )
462 return self._wrap_reduction_result(axis, result)
463
464 def mean(
465 self,
466 *,
467 axis: AxisInt | None = None,
468 dtype: NpDtype | None = None,
469 out=None,
470 keepdims: bool = False,
471 skipna: bool = True,
472 ):
473 nv.validate_mean((), {"dtype": dtype, "out": out, "keepdims": keepdims})
474 result = nanops.nanmean(self._ndarray, axis=axis, skipna=skipna)
475 return self._wrap_reduction_result(axis, result)
476
477 def median(
478 self,
479 *,
480 axis: AxisInt | None = None,
481 out=None,
482 overwrite_input: bool = False,
483 keepdims: bool = False,
484 skipna: bool = True,
485 ):
486 nv.validate_median(
487 (), {"out": out, "overwrite_input": overwrite_input, "keepdims": keepdims}
488 )
489 result = nanops.nanmedian(self._ndarray, axis=axis, skipna=skipna)
490 return self._wrap_reduction_result(axis, result)
491
492 def std(
493 self,
494 *,
495 axis: AxisInt | None = None,
496 dtype: NpDtype | None = None,
497 out=None,
498 ddof: int = 1,
499 keepdims: bool = False,
500 skipna: bool = True,
501 ):
502 nv.validate_stat_ddof_func(
503 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="std"
504 )
505 result = nanops.nanstd(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
506 return self._wrap_reduction_result(axis, result)
507
508 def var(
509 self,
510 *,
511 axis: AxisInt | None = None,
512 dtype: NpDtype | None = None,
513 out=None,
514 ddof: int = 1,
515 keepdims: bool = False,
516 skipna: bool = True,
517 ):
518 nv.validate_stat_ddof_func(
519 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="var"
520 )
521 result = nanops.nanvar(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
522 return self._wrap_reduction_result(axis, result)
523
524 def sem(
525 self,
526 *,
527 axis: AxisInt | None = None,
528 dtype: NpDtype | None = None,
529 out=None,
530 ddof: int = 1,
531 keepdims: bool = False,
532 skipna: bool = True,
533 ):
534 nv.validate_stat_ddof_func(
535 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="sem"
536 )
537 result = nanops.nansem(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
538 return self._wrap_reduction_result(axis, result)
539
540 def kurt(
541 self,
542 *,
543 axis: AxisInt | None = None,
544 dtype: NpDtype | None = None,
545 out=None,
546 keepdims: bool = False,
547 skipna: bool = True,
548 ):
549 nv.validate_stat_ddof_func(
550 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="kurt"
551 )
552 result = nanops.nankurt(self._ndarray, axis=axis, skipna=skipna)
553 return self._wrap_reduction_result(axis, result)
554
555 def skew(
556 self,
557 *,
558 axis: AxisInt | None = None,
559 dtype: NpDtype | None = None,
560 out=None,
561 keepdims: bool = False,
562 skipna: bool = True,
563 ):
564 nv.validate_stat_ddof_func(
565 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="skew"
566 )
567 result = nanops.nanskew(self._ndarray, axis=axis, skipna=skipna)
568 return self._wrap_reduction_result(axis, result)
569
570 # ------------------------------------------------------------------------
571 # Additional Methods
572
573 def to_numpy(
574 self,
575 dtype: npt.DTypeLike | None = None,
576 copy: bool = False,
577 na_value: object = lib.no_default,
578 ) -> np.ndarray:
579 mask = self.isna()
580 if na_value is not lib.no_default and mask.any():
581 result = self._ndarray.copy()
582 result[mask] = na_value
583 else:
584 result = self._ndarray
585 if not copy and self._readonly:
586 result = result.view()
587 result.flags.writeable = False
588
589 result = np.asarray(result, dtype=dtype)
590
591 if copy and result is self._ndarray:
592 result = result.copy()
593
594 return result
595
596 # ------------------------------------------------------------------------
597 # Ops
598
599 def __invert__(self) -> NumpyExtensionArray:
600 return type(self)(~self._ndarray)
601
602 def __neg__(self) -> NumpyExtensionArray:
603 return type(self)(-self._ndarray)
604
605 def __pos__(self) -> NumpyExtensionArray:
606 return type(self)(+self._ndarray)
607
608 def __abs__(self) -> NumpyExtensionArray:
609 return type(self)(abs(self._ndarray))
610
611 def _cmp_method(self, other, op):
612 if isinstance(other, NumpyExtensionArray):
613 other = other._ndarray
614
615 other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
616 pd_op = ops.get_array_op(op)
617 other = ensure_wrapped_if_datetimelike(other)
618 result = pd_op(self._ndarray, other)
619
620 if op is divmod or op is ops.rdivmod:
621 a, b = result
622 if isinstance(a, np.ndarray):
623 # for e.g. op vs TimedeltaArray, we may already
624 # have an ExtensionArray, in which case we do not wrap
625 return self._wrap_ndarray_result(a), self._wrap_ndarray_result(b)
626 return a, b
627
628 if isinstance(result, np.ndarray):
629 # for e.g. multiplication vs TimedeltaArray, we may already
630 # have an ExtensionArray, in which case we do not wrap
631 return self._wrap_ndarray_result(result)
632 return result
633
634 _arith_method = _cmp_method
635
636 def _wrap_ndarray_result(self, result: np.ndarray):
637 # If we have timedelta64[ns] result, return a TimedeltaArray instead
638 # of a NumpyExtensionArray
639 if result.dtype.kind == "m" and is_supported_dtype(result.dtype):
640 from pandas.core.arrays import TimedeltaArray
641
642 return TimedeltaArray._simple_new(result, dtype=result.dtype)
643 return type(self)(result)
644
645 def _formatter(self, boxed: bool = False) -> Callable[[Any], str | None]:
646 # NEP 51: https://github.com/numpy/numpy/pull/22449
647 if self.dtype.kind in "SU":
648 return "'{}'".format
649 elif self.dtype == "object":
650 return repr
651 else:
652 return str