Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from functools import partial
4import operator
5from typing import (
6 TYPE_CHECKING,
7 Any,
8 Literal,
9 Self,
10 cast,
11)
12import warnings
14import numpy as np
16from pandas._config import (
17 get_option,
18 using_string_dtype,
19)
21from pandas._libs import (
22 lib,
23 missing as libmissing,
24)
25from pandas._libs.arrays import NDArrayBacked
26from pandas._libs.lib import ensure_string_array
27from pandas.compat import (
28 HAS_PYARROW,
29 PYARROW_MIN_VERSION,
30)
31from pandas.compat.numpy import function as nv
32from pandas.errors import Pandas4Warning
33from pandas.util._decorators import (
34 set_module,
35)
36from pandas.util._exceptions import find_stack_level
38from pandas.core.dtypes.base import (
39 ExtensionDtype,
40 StorageExtensionDtype,
41 register_extension_dtype,
42)
43from pandas.core.dtypes.common import (
44 is_array_like,
45 is_bool_dtype,
46 is_integer_dtype,
47 is_object_dtype,
48 is_string_dtype,
49 pandas_dtype,
50)
52from pandas.core import (
53 missing,
54 nanops,
55 ops,
56 roperator,
57)
58from pandas.core.algorithms import isin
59from pandas.core.array_algos import masked_reductions
60from pandas.core.arrays.base import ExtensionArray
61from pandas.core.arrays.floating import (
62 FloatingArray,
63 FloatingDtype,
64)
65from pandas.core.arrays.integer import (
66 IntegerArray,
67 IntegerDtype,
68)
69from pandas.core.arrays.numpy_ import NumpyExtensionArray
70from pandas.core.construction import extract_array
71from pandas.core.indexers import check_array_indexer
72from pandas.core.missing import isna
74from pandas.io.formats import printing
76if HAS_PYARROW:
77 import pyarrow as pa
78 import pyarrow.compute as pc
80if TYPE_CHECKING:
81 from collections.abc import MutableMapping
83 import pyarrow
85 from pandas._typing import (
86 ArrayLike,
87 AxisInt,
88 Dtype,
89 DtypeObj,
90 NumpySorter,
91 NumpyValueArrayLike,
92 Scalar,
93 npt,
94 type_t,
95 )
97 from pandas import Series
100@set_module("pandas")
101@register_extension_dtype
102class StringDtype(StorageExtensionDtype):
103 """
104 Extension dtype for string data.
106 .. warning::
108 StringDtype is considered experimental. The implementation and
109 parts of the API may change without warning.
111 Parameters
112 ----------
113 storage : {"python", "pyarrow"}, optional
114 If not given, the value of ``pd.options.mode.string_storage``.
115 na_value : {np.nan, pd.NA}, default pd.NA
116 Whether the dtype follows NaN or NA missing value semantics.
118 Attributes
119 ----------
120 storage
121 na_value
123 Methods
124 -------
125 None
127 See Also
128 --------
129 BooleanDtype : Extension dtype for boolean data.
131 Examples
132 --------
133 >>> pd.StringDtype()
134 <StringDtype(na_value=<NA>)>
136 >>> pd.StringDtype(storage="python")
137 <StringDtype(storage='python', na_value=<NA>)>
138 """
140 @property
141 def name(self) -> str: # type: ignore[override]
142 if self._na_value is libmissing.NA:
143 return "string"
144 else:
145 return "str"
147 #: StringDtype().na_value uses pandas.NA except the implementation that
148 # follows NumPy semantics, which uses nan.
149 @property
150 def na_value(self) -> libmissing.NAType | float: # type: ignore[override]
151 """
152 The missing value representation for this dtype.
154 This value indicates which missing value semantics are used by this dtype.
155 Returns ``np.nan`` for the default string dtype with NumPy semantics,
156 and ``pd.NA`` for the opt-in string dtype with pandas NA semantics.
158 See Also
159 --------
160 isna : Detect missing values.
161 NA : Missing value indicator for nullable dtypes.
163 Examples
164 --------
165 >>> ser = pd.Series(["a", "b"])
166 >>> ser.dtype
167 <StringDtype(na_value=nan)>
168 >>> ser.dtype.na_value
169 nan
170 """
171 return self._na_value
173 @property
174 def storage(self) -> str:
175 """
176 The storage backend for this dtype.
178 Can be either "pyarrow" or "python".
180 See Also
181 --------
182 StringDtype.na_value : The missing value for this dtype.
184 Examples
185 --------
186 >>> ser = pd.Series(["a", "b"])
187 >>> ser.dtype
188 <StringDtype(na_value=nan)>
189 >>> ser.dtype.storage
190 'pyarrow'
191 """
192 return self._storage
194 _metadata = ("storage", "_na_value") # type: ignore[assignment]
196 def __init__(
197 self,
198 storage: str | None = None,
199 na_value: libmissing.NAType | float = libmissing.NA,
200 ) -> None:
201 # infer defaults
202 if storage is None:
203 storage = get_option("mode.string_storage")
204 if storage == "auto":
205 if HAS_PYARROW:
206 storage = "pyarrow"
207 else:
208 storage = "python"
210 # validate options
211 if storage not in {"python", "pyarrow"}:
212 raise ValueError(
213 f"Storage must be 'python' or 'pyarrow'. Got {storage} instead."
214 )
215 if storage == "pyarrow" and not HAS_PYARROW:
216 raise ImportError(
217 f"pyarrow>={PYARROW_MIN_VERSION} is required for PyArrow "
218 "backed StringArray."
219 )
221 if isinstance(na_value, float) and np.isnan(na_value):
222 # when passed a NaN value, always set to np.nan to ensure we use
223 # a consistent NaN value (and we can use `dtype.na_value is np.nan`)
224 na_value = np.nan
225 elif na_value is not libmissing.NA:
226 raise ValueError(f"'na_value' must be np.nan or pd.NA, got {na_value}")
228 self._storage = cast(str, storage)
229 self._na_value = na_value
231 def __repr__(self) -> str:
232 storage = "" if self.storage == "pyarrow" else "storage='python', "
233 return f"<StringDtype({storage}na_value={self._na_value})>"
235 def __eq__(self, other: object) -> bool:
236 # we need to override the base class __eq__ because na_value (NA or NaN)
237 # cannot be checked with normal `==`
238 if isinstance(other, str):
239 # TODO should dtype == "string" work for the NaN variant?
240 if other == "string" or other == self.name: # noqa: PLR1714 (repeated-equality-comparison)
241 return True
242 try:
243 other = self.construct_from_string(other)
244 except (TypeError, ImportError):
245 # TypeError if `other` is not a valid string for StringDtype
246 # ImportError if pyarrow is not installed for "string[pyarrow]"
247 return False
248 if isinstance(other, type(self)):
249 return self.storage == other.storage and self.na_value is other.na_value
250 return False
252 def __setstate__(self, state: MutableMapping[str, Any]) -> None:
253 # back-compat for pandas < 2.3, where na_value did not yet exist
254 self._storage = state.pop("storage", "python")
255 self._na_value = state.pop("_na_value", libmissing.NA)
257 def __hash__(self) -> int:
258 # need to override __hash__ as well because of overriding __eq__
259 return super().__hash__()
261 def __reduce__(self):
262 return StringDtype, (self.storage, self.na_value)
264 @property
265 def type(self) -> type[str]:
266 return str
268 @classmethod
269 def construct_from_string(cls, string) -> Self:
270 """
271 Construct a StringDtype from a string.
273 Parameters
274 ----------
275 string : str
276 The type of the name. The storage type will be taking from `string`.
277 Valid options and their storage types are
279 ========================== ==============================================
280 string result storage
281 ========================== ==============================================
282 ``'string'`` pd.options.mode.string_storage, default python
283 ``'string[python]'`` python
284 ``'string[pyarrow]'`` pyarrow
285 ========================== ==============================================
287 Returns
288 -------
289 StringDtype
291 Raise
292 -----
293 TypeError
294 If the string is not a valid option.
295 """
296 if not isinstance(string, str):
297 raise TypeError(
298 f"'construct_from_string' expects a string, got {type(string)}"
299 )
300 if string == "string":
301 return cls()
302 elif string == "str" and using_string_dtype():
303 return cls(na_value=np.nan)
304 elif string == "string[python]":
305 return cls(storage="python")
306 elif string == "string[pyarrow]":
307 return cls(storage="pyarrow")
308 else:
309 raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
311 def construct_array_type(self) -> type_t[BaseStringArray]:
312 """
313 Return the array type associated with this dtype.
315 Returns
316 -------
317 type
318 """
319 from pandas.core.arrays.string_arrow import (
320 ArrowStringArray,
321 )
323 if self.storage == "python" and self._na_value is libmissing.NA:
324 return StringArray
325 elif self.storage == "pyarrow" and self._na_value is libmissing.NA:
326 return ArrowStringArray
327 elif self.storage == "python":
328 return StringArray
329 else:
330 return ArrowStringArray
332 def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
333 storages = set()
334 na_values = set()
336 for dtype in dtypes:
337 if isinstance(dtype, StringDtype):
338 storages.add(dtype.storage)
339 na_values.add(dtype.na_value)
340 elif isinstance(dtype, np.dtype) and dtype.kind in ("U", "T"):
341 continue
342 else:
343 return None
345 if len(storages) == 2:
346 # if both python and pyarrow storage -> priority to pyarrow
347 storage = "pyarrow"
348 else:
349 storage = next(iter(storages))
351 na_value: libmissing.NAType | float
352 if len(na_values) == 2:
353 # if both NaN and NA -> priority to NA
354 na_value = libmissing.NA
355 else:
356 na_value = next(iter(na_values))
358 return StringDtype(storage=storage, na_value=na_value)
360 def __from_arrow__(
361 self, array: pyarrow.Array | pyarrow.ChunkedArray
362 ) -> BaseStringArray:
363 """
364 Construct StringArray from pyarrow Array/ChunkedArray.
365 """
366 if self.storage == "pyarrow":
367 from pandas.core.arrays.string_arrow import (
368 ArrowStringArray,
369 _check_pyarrow_available,
370 )
372 _check_pyarrow_available()
374 if not pa.types.is_large_string(array.type):
375 array = pc.cast(array, pa.large_string())
377 return ArrowStringArray(array, dtype=self)
379 else:
380 import pyarrow
382 if isinstance(array, pyarrow.Array):
383 chunks = [array]
384 else:
385 # pyarrow.ChunkedArray
386 chunks = array.chunks
388 results = []
389 for arr in chunks:
390 # convert chunk by chunk to numpy and concatenate then, to avoid
391 # overflow for large string data when concatenating the pyarrow arrays
392 arr = arr.to_numpy(zero_copy_only=False)
393 arr = ensure_string_array(arr, na_value=self.na_value)
394 results.append(arr)
396 if len(chunks) == 0:
397 arr = np.array([], dtype=object)
398 else:
399 arr = np.concatenate(results)
401 # Bypass validation inside StringArray constructor, see GH#47781
402 new_string_array = StringArray.__new__(StringArray)
403 NDArrayBacked.__init__(new_string_array, arr, self)
404 return new_string_array
407class BaseStringArray(ExtensionArray):
408 """
409 Mixin class for StringArray, ArrowStringArray.
410 """
412 dtype: StringDtype
414 # TODO(4.0): Once the deprecation here is enforced, this method can be
415 # removed and we use the parent class method instead.
416 def _logical_method(self, other, op):
417 if (
418 op in (roperator.ror_, roperator.rand_, roperator.rxor)
419 and isinstance(other, np.ndarray)
420 and other.dtype == bool
421 ):
422 # GH#60234 backward compatibility for the move to StringDtype in 3.0
423 op_name = op.__name__[1:].strip("_")
424 warnings.warn(
425 f"'{op_name}' operations between boolean dtype and {self.dtype} are "
426 "deprecated and will raise in a future version. Explicitly "
427 "cast the strings to a boolean dtype before operating instead.",
428 Pandas4Warning,
429 stacklevel=find_stack_level(),
430 )
431 return op(other, self.astype(bool))
432 return NotImplemented
434 def tolist(self) -> list:
435 """
436 Return a list of the value.
438 These are each a scalar type, which is a Python scalar
439 (for str, int, float) or pandas scalar
440 (for Timestamp/Timedelta/Interval/Period)
442 Returns
443 ----------
444 list
446 Examples
447 ----------
448 >>> arr = pd.array(["a", "b", "c"])
449 >>> arr.tolist()
450 ['a', 'b', 'c']
451 """
452 if self.ndim > 1:
453 return [x.tolist() for x in self]
454 return list(self.to_numpy())
456 def _formatter(self, boxed: bool = False):
457 formatter = partial(
458 printing.pprint_thing,
459 escape_chars=("\t", "\r", "\n"),
460 quote_strings=not boxed,
461 )
462 return formatter
464 def _str_map(
465 self,
466 f,
467 na_value=lib.no_default,
468 dtype: Dtype | None = None,
469 convert: bool = True,
470 ):
471 if self.dtype.na_value is np.nan:
472 return self._str_map_nan_semantics(f, na_value=na_value, dtype=dtype)
474 from pandas.arrays import BooleanArray
476 if dtype is None:
477 dtype = self.dtype
478 if na_value is lib.no_default:
479 na_value = self.dtype.na_value
481 mask = isna(self)
482 arr = np.asarray(self)
484 if is_integer_dtype(dtype) or is_bool_dtype(dtype):
485 constructor: type[IntegerArray | BooleanArray]
486 if is_integer_dtype(dtype):
487 constructor = IntegerArray
488 else:
489 constructor = BooleanArray
491 na_value_is_na = isna(na_value)
492 if na_value_is_na:
493 na_value = 1
494 elif dtype == np.dtype("bool"):
495 # GH#55736
496 na_value = bool(na_value)
497 result = lib.map_infer_mask(
498 arr,
499 f,
500 mask.view("uint8"),
501 convert=False,
502 na_value=na_value,
503 # error: Argument 1 to "dtype" has incompatible type
504 # "Union[ExtensionDtype, str, dtype[Any], Type[object]]"; expected
505 # "Type[object]"
506 dtype=np.dtype(cast(type, dtype)),
507 )
509 if not na_value_is_na:
510 mask[:] = False
512 return constructor(result, mask)
514 else:
515 return self._str_map_str_or_object(dtype, na_value, arr, f, mask)
517 def _str_map_str_or_object(
518 self,
519 dtype,
520 na_value,
521 arr: np.ndarray,
522 f,
523 mask: npt.NDArray[np.bool_],
524 ):
525 # _str_map helper for case where dtype is either string dtype or object
526 if is_string_dtype(dtype) and not is_object_dtype(dtype):
527 # i.e. StringDtype
528 result = lib.map_infer_mask(
529 arr, f, mask.view("uint8"), convert=False, na_value=na_value
530 )
531 if self.dtype.storage == "pyarrow":
532 import pyarrow as pa
534 # TODO: shouldn't this already be caught my passed mask?
535 # it isn't in test_extract_expand_capture_groups_index
536 # mask = mask | np.array(
537 # [x is libmissing.NA for x in result], dtype=bool
538 # )
540 result = pa.array(
541 result, mask=mask, type=pa.large_string(), from_pandas=True
542 )
543 # error: "BaseStringArray" has no attribute "_from_pyarrow_array"
544 return self._from_pyarrow_array(result) # type: ignore[attr-defined]
545 else:
546 # StringArray
547 # error: Too many arguments for "BaseStringArray"
548 return type(self)(result, dtype=self.dtype) # type: ignore[call-arg]
550 else:
551 # This is when the result type is object. We reach this when
552 # -> We know the result type is truly object (e.g. .encode returns bytes
553 # or .findall returns a list).
554 # -> We don't know the result type. E.g. `.get` can return anything.
555 return lib.map_infer_mask(arr, f, mask.view("uint8"))
557 def _str_map_nan_semantics(
558 self, f, na_value=lib.no_default, dtype: Dtype | None = None
559 ):
560 if dtype is None:
561 dtype = self.dtype
562 if na_value is lib.no_default:
563 if is_bool_dtype(dtype):
564 # NaN propagates as False
565 na_value = False
566 else:
567 na_value = self.dtype.na_value
569 mask = isna(self)
570 arr = np.asarray(self)
572 if is_integer_dtype(dtype) or is_bool_dtype(dtype):
573 na_value_is_na = isna(na_value)
574 if na_value_is_na:
575 if is_integer_dtype(dtype):
576 na_value = 0
577 else:
578 # NaN propagates as False
579 na_value = False
581 result = lib.map_infer_mask(
582 arr,
583 f,
584 mask.view("uint8"),
585 convert=False,
586 na_value=na_value,
587 dtype=np.dtype(cast(type, dtype)),
588 )
589 if na_value_is_na and is_integer_dtype(dtype) and mask.any():
590 # TODO: we could alternatively do this check before map_infer_mask
591 # and adjust the dtype/na_value we pass there. Which is more
592 # performant?
593 result = result.astype("float64")
594 result[mask] = np.nan
596 return result
598 else:
599 return self._str_map_str_or_object(dtype, na_value, arr, f, mask)
601 def view(self, dtype: Dtype | None = None) -> Self:
602 if dtype is not None:
603 raise TypeError("Cannot change data-type for string array.")
604 return super().view()
607@set_module("pandas.arrays")
608# error: Definition of "_concat_same_type" in base class "NDArrayBacked" is
609# incompatible with definition in base class "ExtensionArray"
610class StringArray(BaseStringArray, NumpyExtensionArray): # type: ignore[misc]
611 """
612 Extension array for string data.
614 .. warning::
616 StringArray is considered experimental. The implementation and
617 parts of the API may change without warning.
619 Parameters
620 ----------
621 values : array-like
622 The array of data.
624 .. warning::
626 Currently, this expects an object-dtype ndarray
627 where the elements are Python strings
628 or nan-likes (``None``, ``np.nan``, ``NA``).
629 This may change without warning in the future. Use
630 :meth:`pandas.array` with ``dtype="string"`` for a stable way of
631 creating a `StringArray` from any sequence.
633 StringArray accepts array-likes containing
634 nan-likes(``None``, ``np.nan``) for the ``values`` parameter
635 in addition to strings and :attr:`pandas.NA`
637 dtype : StringDtype
638 Dtype for the array.
639 copy : bool, default False
640 Whether to copy the array of data.
642 Attributes
643 ----------
644 None
646 Methods
647 -------
648 None
650 See Also
651 --------
652 :func:`array`
653 The recommended function for creating a StringArray.
654 Series.str
655 The string methods are available on Series backed by
656 a StringArray.
658 Notes
659 -----
660 StringArray returns a BooleanArray for comparison methods.
662 Examples
663 --------
664 >>> pd.array(["This is", "some text", None, "data."], dtype="string")
665 <ArrowStringArray>
666 ['This is', 'some text', <NA>, 'data.']
667 Length: 4, dtype: string
669 Unlike arrays instantiated with ``dtype="object"``, ``StringArray``
670 will convert the values to strings.
672 >>> pd.array(["1", 1], dtype="object")
673 <NumpyExtensionArray>
674 ['1', 1]
675 Length: 2, dtype: object
676 >>> pd.array(["1", 1], dtype="string")
677 <ArrowStringArray>
678 ['1', '1']
679 Length: 2, dtype: string
681 However, instantiating StringArrays directly with non-strings will raise an error.
683 For comparison methods, `StringArray` returns a :class:`pandas.BooleanArray`:
685 >>> pd.array(["a", None, "c"], dtype="string[python]") == "a"
686 <BooleanArray>
687 [True, <NA>, False]
688 Length: 3, dtype: boolean
689 """
691 # undo the NumpyExtensionArray hack
692 _typ = "extension"
694 def __init__(
695 self, values, *, dtype: StringDtype | None = None, copy: bool = False
696 ) -> None:
697 if dtype is None:
698 dtype = StringDtype()
699 values = extract_array(values)
701 super().__init__(values, copy=copy)
702 if not isinstance(values, type(self)):
703 self._validate(dtype)
704 NDArrayBacked.__init__(
705 self,
706 self._ndarray,
707 dtype,
708 )
710 def _validate(self, dtype: StringDtype) -> None:
711 """Validate that we only store NA or strings."""
713 if dtype._na_value is libmissing.NA:
714 if len(self._ndarray) and not lib.is_string_array(
715 self._ndarray, skipna=True
716 ):
717 raise ValueError(
718 "StringArray requires a sequence of strings or pandas.NA"
719 )
720 if self._ndarray.dtype != "object":
721 raise ValueError(
722 "StringArray requires a sequence of strings or pandas.NA. Got "
723 f"'{self._ndarray.dtype}' dtype instead."
724 )
725 # Check to see if need to convert Na values to pd.NA
726 if self._ndarray.ndim > 2:
727 # Ravel if ndims > 2 b/c no cythonized version available
728 lib.convert_nans_to_NA(self._ndarray.ravel("K"))
729 else:
730 lib.convert_nans_to_NA(self._ndarray)
731 else:
732 # Validate that we only store NaN or strings.
733 if len(self._ndarray) and not lib.is_string_array(
734 self._ndarray, skipna=True
735 ):
736 raise ValueError("StringArray requires a sequence of strings or NaN")
737 if self._ndarray.dtype != "object":
738 raise ValueError(
739 "StringArray requires a sequence of strings "
740 "or NaN. Got '{self._ndarray.dtype}' dtype instead."
741 )
742 # TODO validate or force NA/None to NaN
744 def _validate_scalar(self, value):
745 # used by NDArrayBackedExtensionIndex.insert
746 if isna(value):
747 return self.dtype.na_value
748 elif not isinstance(value, str):
749 raise TypeError(
750 f"Invalid value '{value}' for dtype '{self.dtype}'. Value should be a "
751 f"string or missing value, got '{type(value).__name__}' instead."
752 )
753 return value
755 @classmethod
756 def _from_sequence(
757 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
758 ) -> Self:
759 if dtype and not (isinstance(dtype, str) and dtype == "string"):
760 dtype = pandas_dtype(dtype)
761 assert isinstance(dtype, StringDtype) and dtype.storage == "python"
762 elif using_string_dtype():
763 dtype = StringDtype(storage="python", na_value=np.nan)
764 else:
765 dtype = StringDtype(storage="python")
767 from pandas.core.arrays.masked import BaseMaskedArray
769 na_value = dtype.na_value
770 if isinstance(scalars, BaseMaskedArray):
771 # avoid costly conversion to object dtype
772 na_values = scalars._mask
773 result = scalars._data
774 result = lib.ensure_string_array(
775 result, copy=copy, convert_na_value=False, skipna=False
776 )
777 result[na_values] = na_value
779 else:
780 if lib.is_pyarrow_array(scalars):
781 # pyarrow array; we cannot rely on the "to_numpy" check in
782 # ensure_string_array because calling scalars.to_numpy would set
783 # zero_copy_only to True which caused problems see GH#52076
784 scalars = np.array(scalars)
785 # convert non-na-likes to str, and nan-likes to StringDtype().na_value
786 result = lib.ensure_string_array(scalars, na_value=na_value, copy=copy)
788 # Manually creating new array avoids the validation step in the __init__, so is
789 # faster. Refactor need for validation?
790 new_string_array = cls.__new__(cls)
791 NDArrayBacked.__init__(new_string_array, result, dtype)
793 return new_string_array
795 @classmethod
796 def _from_sequence_of_strings(
797 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
798 ) -> Self:
799 return cls._from_sequence(strings, dtype=dtype, copy=copy)
801 def _cast_pointwise_result(self, values) -> ArrayLike:
802 result = super()._cast_pointwise_result(values)
803 if isinstance(result.dtype, StringDtype):
804 # Ensure we retain our same na_value/storage
805 result = result.astype(self.dtype)
806 return result
808 @classmethod
809 def _empty(cls, shape, dtype) -> StringArray:
810 values = np.empty(shape, dtype=object)
811 values[:] = dtype.na_value
812 return cls(values, dtype=dtype).astype(dtype, copy=False)
814 def __arrow_array__(self, type=None):
815 """
816 Convert myself into a pyarrow Array.
817 """
818 import pyarrow as pa
820 if type is None:
821 type = pa.string()
823 values = self._ndarray.copy()
824 values[self.isna()] = None
825 return pa.array(values, type=type)
827 def _values_for_factorize(self) -> tuple[np.ndarray, libmissing.NAType | float]: # type: ignore[override]
828 arr = self._ndarray
830 return arr, self.dtype.na_value
832 def _maybe_convert_setitem_value(self, value):
833 """Maybe convert value to be StringArray compatible."""
834 if lib.is_scalar(value):
835 if isna(value):
836 value = self.dtype.na_value
837 elif not isinstance(value, str):
838 raise TypeError(
839 f"Invalid value '{value}' for dtype '{self.dtype}'. Value should "
840 f"be a string or missing value, got '{type(value).__name__}' "
841 "instead."
842 )
843 else:
844 value = extract_array(value, extract_numpy=True)
845 if not is_array_like(value):
846 value = np.asarray(value, dtype=object)
847 elif isinstance(value.dtype, type(self.dtype)):
848 return value
849 else:
850 # cast categories and friends to arrays to see if values are
851 # compatible, compatibility with arrow backed strings
852 value = np.asarray(value)
853 if len(value) and not lib.is_string_array(value, skipna=True):
854 raise TypeError(
855 "Invalid value for dtype 'str'. Value should be a "
856 "string or missing value (or array of those)."
857 )
858 return value
860 def __setitem__(self, key, value) -> None:
861 if self._readonly:
862 raise ValueError("Cannot modify read-only array")
864 value = self._maybe_convert_setitem_value(value)
866 key = check_array_indexer(self, key)
867 scalar_key = lib.is_scalar(key)
868 scalar_value = lib.is_scalar(value)
869 if scalar_key and not scalar_value:
870 raise ValueError("setting an array element with a sequence.")
872 if not scalar_value:
873 if value.dtype == self.dtype:
874 value = value._ndarray
875 else:
876 value = np.asarray(value)
877 mask = isna(value)
878 if mask.any():
879 value = value.copy()
880 value[isna(value)] = self.dtype.na_value
882 super().__setitem__(key, value)
884 def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
885 # the super() method NDArrayBackedExtensionArray._putmask uses
886 # np.putmask which doesn't properly handle None/pd.NA, so using the
887 # base class implementation that uses __setitem__
888 ExtensionArray._putmask(self, mask, value)
890 def _where(self, mask: npt.NDArray[np.bool_], value) -> Self:
891 # the super() method NDArrayBackedExtensionArray._where uses
892 # np.putmask which doesn't properly handle None/pd.NA, so using the
893 # base class implementation that uses __setitem__
894 return ExtensionArray._where(self, mask, value)
896 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
897 if isinstance(values, BaseStringArray) or (
898 isinstance(values, ExtensionArray) and is_string_dtype(values.dtype)
899 ):
900 values = values.astype(self.dtype, copy=False)
901 else:
902 if not lib.is_string_array(np.asarray(values), skipna=True):
903 values = np.array(
904 [val for val in values if isinstance(val, str) or isna(val)],
905 dtype=object,
906 )
907 if not len(values):
908 return np.zeros(self.shape, dtype=bool)
910 values = self._from_sequence(values, dtype=self.dtype)
912 return isin(np.asarray(self), np.asarray(values))
914 def astype(self, dtype, copy: bool = True):
915 dtype = pandas_dtype(dtype)
917 if dtype == self.dtype:
918 if copy:
919 return self.copy()
920 return self
922 elif isinstance(dtype, IntegerDtype):
923 arr = self._ndarray.copy()
924 mask = self.isna()
925 arr[mask] = 0
926 values = arr.astype(dtype.numpy_dtype)
927 return IntegerArray(values, mask, copy=False)
928 elif isinstance(dtype, FloatingDtype):
929 arr_ea = self.copy()
930 mask = self.isna()
931 arr_ea[mask] = "0"
932 values = arr_ea.astype(dtype.numpy_dtype)
933 return FloatingArray(values, mask, copy=False)
934 elif isinstance(dtype, ExtensionDtype):
935 # Skip the NumpyExtensionArray.astype method
936 return ExtensionArray.astype(self, dtype, copy)
937 elif np.issubdtype(dtype, np.floating):
938 arr = self._ndarray.copy()
939 mask = self.isna()
940 arr[mask] = 0
941 values = arr.astype(dtype)
942 values[mask] = np.nan
943 return values
945 return super().astype(dtype, copy)
947 def _reduce(
948 self,
949 name: str,
950 *,
951 skipna: bool = True,
952 keepdims: bool = False,
953 axis: AxisInt | None = 0,
954 **kwargs,
955 ):
956 if self.dtype.na_value is np.nan and name in ["any", "all"]:
957 if name == "any":
958 return nanops.nanany(self._ndarray, skipna=skipna)
959 else:
960 return nanops.nanall(self._ndarray, skipna=skipna)
962 if name in ["min", "max", "argmin", "argmax", "sum"]:
963 result = getattr(self, name)(skipna=skipna, axis=axis, **kwargs)
964 if keepdims:
965 return self._from_sequence([result], dtype=self.dtype)
966 return result
968 raise TypeError(f"Cannot perform reduction '{name}' with string dtype")
970 def _accumulate(self, name: str, *, skipna: bool = True, **kwargs) -> StringArray:
971 """
972 Return an ExtensionArray performing an accumulation operation.
974 The underlying data type might change.
976 Parameters
977 ----------
978 name : str
979 Name of the function, supported values are:
980 - cummin
981 - cummax
982 - cumsum
983 - cumprod
984 skipna : bool, default True
985 If True, skip NA values.
986 **kwargs
987 Additional keyword arguments passed to the accumulation function.
988 Currently, there is no supported kwarg.
990 Returns
991 -------
992 array
994 Raises
995 ------
996 NotImplementedError : subclass does not define accumulations
997 """
998 if name == "cumprod":
999 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
1000 raise TypeError(msg)
1002 # We may need to strip out trailing NA values
1003 tail: np.ndarray | None = None
1004 na_mask: np.ndarray | None = None
1005 ndarray = self._ndarray
1006 np_func = {
1007 "cumsum": np.cumsum,
1008 "cummin": np.minimum.accumulate,
1009 "cummax": np.maximum.accumulate,
1010 }[name]
1012 if self._hasna:
1013 na_mask = cast("npt.NDArray[np.bool_]", isna(ndarray))
1014 if np.all(na_mask):
1015 return type(self)(ndarray, dtype=self.dtype)
1016 if skipna:
1017 if name == "cumsum":
1018 ndarray = np.where(na_mask, "", ndarray)
1019 else:
1020 # We can retain the running min/max by forward/backward filling.
1021 ndarray = ndarray.copy()
1022 missing.pad_or_backfill_inplace(
1023 ndarray,
1024 method="pad",
1025 axis=0,
1026 )
1027 missing.pad_or_backfill_inplace(
1028 ndarray,
1029 method="backfill",
1030 axis=0,
1031 )
1032 else:
1033 # When not skipping NA values, the result should be null from
1034 # the first NA value onward.
1035 idx = np.argmax(na_mask)
1036 tail = np.empty(len(ndarray) - idx, dtype="object")
1037 tail[:] = self.dtype.na_value
1038 ndarray = ndarray[:idx]
1040 # mypy: Cannot call function of unknown type
1041 np_result = np_func(ndarray) # type: ignore[operator]
1043 if tail is not None:
1044 np_result = np.hstack((np_result, tail))
1045 elif na_mask is not None:
1046 # Argument 2 to "where" has incompatible type "NAType | float"
1047 np_result = np.where(na_mask, self.dtype.na_value, np_result) # type: ignore[arg-type]
1049 result = type(self)(np_result, dtype=self.dtype)
1050 return result
1052 def _wrap_reduction_result(self, axis: AxisInt | None, result) -> Any:
1053 if self.dtype.na_value is np.nan and result is libmissing.NA:
1054 # the masked_reductions use pd.NA -> convert to np.nan
1055 return np.nan
1056 return super()._wrap_reduction_result(axis, result)
1058 def min(self, axis=None, skipna: bool = True, **kwargs) -> Scalar:
1059 nv.validate_min((), kwargs)
1060 result = masked_reductions.min(
1061 values=self.to_numpy(), mask=self.isna(), skipna=skipna
1062 )
1063 return self._wrap_reduction_result(axis, result)
1065 def max(self, axis=None, skipna: bool = True, **kwargs) -> Scalar:
1066 nv.validate_max((), kwargs)
1067 result = masked_reductions.max(
1068 values=self.to_numpy(), mask=self.isna(), skipna=skipna
1069 )
1070 return self._wrap_reduction_result(axis, result)
1072 def sum(
1073 self,
1074 *,
1075 axis: AxisInt | None = None,
1076 skipna: bool = True,
1077 min_count: int = 0,
1078 **kwargs,
1079 ) -> Scalar:
1080 nv.validate_sum((), kwargs)
1081 result = masked_reductions.sum(
1082 values=self._ndarray,
1083 mask=self.isna(),
1084 skipna=skipna,
1085 min_count=min_count,
1086 initial="",
1087 )
1088 return self._wrap_reduction_result(axis, result)
1090 def value_counts(self, dropna: bool = True) -> Series:
1091 result = super().value_counts(dropna=dropna)
1093 if self.dtype.na_value is libmissing.NA:
1094 result = result.astype("Int64")
1095 return result
1097 def memory_usage(self, deep: bool = False) -> int:
1098 result = self._ndarray.nbytes
1099 if deep:
1100 return result + lib.memory_usage_of_objects(self._ndarray)
1101 return result
1103 def searchsorted(
1104 self,
1105 value: NumpyValueArrayLike | ExtensionArray,
1106 side: Literal["left", "right"] = "left",
1107 sorter: NumpySorter | None = None,
1108 ) -> npt.NDArray[np.intp] | np.intp:
1109 """
1110 Find indices where elements should be inserted to maintain order.
1112 Find the indices into a sorted array `self` (a) such that, if the
1113 corresponding elements in `value` were inserted before the indices,
1114 the order of `self` would be preserved.
1116 Assuming that `self` is sorted:
1118 ====== ================================
1119 `side` returned index `i` satisfies
1120 ====== ================================
1121 left ``self[i-1] < value <= self[i]``
1122 right ``self[i-1] <= value < self[i]``
1123 ====== ================================
1125 Parameters
1126 ----------
1127 value : array-like, list or scalar
1128 Value(s) to insert into `self`.
1129 side : {'left', 'right'}, optional
1130 If 'left', the index of the first suitable location found is given.
1131 If 'right', return the last such index. If there is no suitable
1132 index, return either 0 or N (where N is the length of `self`).
1133 sorter : 1-D array-like, optional
1134 Optional array of integer indices that sort array a into ascending
1135 order. They are typically the result of argsort.
1137 Returns
1138 -------
1139 array of ints or int
1140 If value is array-like, array of insertion points.
1141 If value is scalar, a single integer.
1143 See Also
1144 --------
1145 numpy.searchsorted : Similar method from NumPy.
1147 Examples
1148 --------
1149 >>> arr = pd.array([1, 2, 3, 5])
1150 >>> arr.searchsorted([4])
1151 array([3])
1152 """
1154 # GH#65837: avoid O(n) scan; NA confined to array ends in sorted data.
1155 # When sorter is given, the sorted order is ndarray[sorter], so check
1156 # the first/last positions via sorter instead of raw ndarray positions.
1157 ndarray = self._ndarray
1158 if len(ndarray):
1159 if sorter is None:
1160 has_na = libmissing.checknull(ndarray[0]) or libmissing.checknull(
1161 ndarray[-1]
1162 )
1163 else:
1164 has_na = libmissing.checknull(
1165 ndarray[sorter[0]]
1166 ) or libmissing.checknull(ndarray[sorter[-1]])
1167 else:
1168 has_na = False
1169 if has_na:
1170 raise ValueError(
1171 "searchsorted requires array to be sorted, which is impossible "
1172 "with NAs present."
1173 )
1174 return super().searchsorted(value=value, side=side, sorter=sorter)
1176 def _cmp_method(self, other, op):
1177 from pandas.arrays import (
1178 ArrowExtensionArray,
1179 BooleanArray,
1180 )
1182 if (
1183 isinstance(other, BaseStringArray)
1184 and self.dtype.na_value is not libmissing.NA
1185 and other.dtype.na_value is libmissing.NA
1186 ):
1187 # NA has priority of NaN semantics
1188 return op(self.astype(other.dtype, copy=False), other)
1190 if isinstance(other, ArrowExtensionArray):
1191 if isinstance(other, BaseStringArray):
1192 # pyarrow storage has priority over python storage
1193 # (except if we have NA semantics and other not)
1194 if not (
1195 self.dtype.na_value is libmissing.NA
1196 and other.dtype.na_value is not libmissing.NA
1197 ):
1198 return NotImplemented
1199 else:
1200 return NotImplemented
1202 if isinstance(other, StringArray):
1203 other = other._ndarray
1205 mask = isna(self) | isna(other)
1206 valid = ~mask
1208 if lib.is_list_like(other):
1209 if len(other) != len(self):
1210 # prevent improper broadcasting when other is 2D
1211 raise ValueError(
1212 f"Lengths of operands do not match: {len(self)} != {len(other)}"
1213 )
1215 # for array-likes, first filter out NAs before converting to numpy
1216 if not is_array_like(other):
1217 other = np.asarray(other)
1218 other = other[valid]
1220 other_dtype = getattr(other, "dtype", None)
1221 if op.__name__.strip("_") in ["mul", "rmul"] and (
1222 lib.is_bool(other) or lib.is_np_dtype(other_dtype, "b")
1223 ):
1224 # GH#62595
1225 raise TypeError(
1226 "Cannot multiply StringArray by bools. "
1227 "Explicitly cast to integers instead."
1228 )
1230 if op.__name__ in ops.ARITHMETIC_BINOPS:
1231 result = np.empty_like(self._ndarray, dtype="object")
1232 result[mask] = self.dtype.na_value
1233 result[valid] = op(self._ndarray[valid], other)
1234 if not lib.is_string_array(result, skipna=True):
1235 return result
1236 return self._from_backing_data(result)
1237 else:
1238 # logical
1239 result = np.zeros(len(self._ndarray), dtype="bool")
1240 result[valid] = op(self._ndarray[valid], other)
1241 res_arr = BooleanArray(result, mask)
1242 if self.dtype.na_value is np.nan:
1243 if op == operator.ne:
1244 return res_arr.to_numpy(np.bool_, na_value=True)
1245 else:
1246 return res_arr.to_numpy(np.bool_, na_value=False)
1247 return res_arr
1249 _arith_method = _cmp_method
1251 def _str_zfill(self, width: int) -> Self:
1252 return self._str_map(lambda x: x.zfill(width))