1from __future__ import annotations
2
3from functools import partial
4import operator
5from typing import (
6 TYPE_CHECKING,
7 Any,
8 Literal,
9 Self,
10 cast,
11 overload,
12)
13import warnings
14
15import numpy as np
16
17from pandas._config import using_string_dtype
18from pandas._config.config import _global_config as config
19
20from pandas._libs import (
21 lib,
22 missing as libmissing,
23)
24from pandas._libs.arrays import NDArrayBacked
25from pandas._libs.lib import ensure_string_array
26from pandas.compat import (
27 HAS_PYARROW,
28 PYARROW_MIN_VERSION,
29 PYPY,
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
37
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_deprecate_non_pandas,
45 is_bool_dtype,
46 is_integer_dtype,
47 is_object_dtype,
48 is_string_dtype,
49 pandas_dtype,
50)
51
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
73
74from pandas.io.formats import printing
75
76if HAS_PYARROW:
77 import pyarrow as pa
78 import pyarrow.compute as pc
79
80if TYPE_CHECKING:
81 from collections.abc import MutableMapping
82
83 import pyarrow
84
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 )
96
97 from pandas import Series
98
99
100@set_module("pandas")
101@register_extension_dtype
102class StringDtype(StorageExtensionDtype):
103 """
104 Extension dtype for string data.
105
106 .. warning::
107
108 StringDtype is considered experimental. The implementation and
109 parts of the API may change without warning.
110
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.
117
118 Attributes
119 ----------
120 storage
121 na_value
122
123 Methods
124 -------
125 None
126
127 See Also
128 --------
129 BooleanDtype : Extension dtype for boolean data.
130
131 Examples
132 --------
133 >>> pd.StringDtype()
134 <StringDtype(na_value=<NA>)>
135
136 >>> pd.StringDtype(storage="python")
137 <StringDtype(storage='python', na_value=<NA>)>
138 """
139
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"
146
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.
153
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.
157
158 See Also
159 --------
160 isna : Detect missing values.
161 NA : Missing value indicator for nullable dtypes.
162
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
172
173 @property
174 def storage(self) -> str:
175 """
176 The storage backend for this dtype.
177
178 Can be either "pyarrow" or "python".
179
180 See Also
181 --------
182 StringDtype.na_value : The missing value for this dtype.
183
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
193
194 _metadata = ("storage", "_na_value") # type: ignore[assignment]
195
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 = config["mode"]["string_storage"]
204 if storage == "auto":
205 if HAS_PYARROW:
206 storage = "pyarrow"
207 else:
208 storage = "python"
209
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 )
220
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}")
227
228 self._storage = cast("str", storage)
229 self._na_value = na_value
230
231 def __repr__(self) -> str:
232 storage = "" if self.storage == "pyarrow" else "storage='python', "
233 return f"<StringDtype({storage}na_value={self._na_value})>"
234
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
251
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)
256
257 def __hash__(self) -> int:
258 # need to override __hash__ as well because of overriding __eq__
259 return super().__hash__()
260
261 def __reduce__(self):
262 return StringDtype, (self.storage, self.na_value)
263
264 @property
265 def type(self) -> type[str]:
266 return str
267
268 @classmethod
269 def construct_from_string(cls, string) -> Self:
270 """
271 Construct a StringDtype from a string.
272
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
278
279 ========================== ==============================================
280 string result storage
281 ========================== ==============================================
282 ``'string'`` pd.options.mode.string_storage, default python
283 ``'string[python]'`` python
284 ``'string[pyarrow]'`` pyarrow
285 ========================== ==============================================
286
287 Returns
288 -------
289 StringDtype
290
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}'")
310
311 def construct_array_type(self) -> type_t[BaseStringArray]:
312 """
313 Return the array type associated with this dtype.
314
315 Returns
316 -------
317 type
318 """
319 from pandas.core.arrays.string_arrow import (
320 ArrowStringArray,
321 )
322
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
331
332 def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
333 storages = set()
334 na_values = set()
335
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
344
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))
350
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))
357
358 return StringDtype(storage=storage, na_value=na_value)
359
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 )
371
372 _check_pyarrow_available()
373
374 if not pa.types.is_large_string(array.type):
375 array = pc.cast(array, pa.large_string())
376
377 return ArrowStringArray(array, dtype=self)
378
379 else:
380 import pyarrow
381
382 if isinstance(array, pyarrow.Array):
383 chunks = [array]
384 else:
385 # pyarrow.ChunkedArray
386 chunks = array.chunks
387
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)
395
396 if len(chunks) == 0:
397 arr = np.array([], dtype=object)
398 else:
399 arr = np.concatenate(results)
400
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
405
406
407class BaseStringArray(ExtensionArray):
408 """
409 Mixin class for StringArray, ArrowStringArray.
410 """
411
412 dtype: StringDtype
413
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 # the GH#60234 arm below silently broadcasts a 2-D ndarray
418 ops.raise_if_2d(other)
419
420 if (
421 op in (roperator.ror_, roperator.rand_, roperator.rxor)
422 and isinstance(other, np.ndarray)
423 and other.dtype == bool
424 ):
425 # GH#60234 backward compatibility for the move to StringDtype in 3.0
426 op_name = op.__name__[1:].strip("_")
427 warnings.warn(
428 f"'{op_name}' operations between boolean dtype and {self.dtype} are "
429 "deprecated and will raise in a future version. Explicitly "
430 "cast the strings to a boolean dtype before operating instead.",
431 Pandas4Warning,
432 stacklevel=find_stack_level(),
433 )
434 return op(other, self.astype(bool))
435 return NotImplemented
436
437 def tolist(self) -> list[Any]:
438 """
439 Return a list of the value.
440
441 These are each a scalar type, which is a Python scalar
442 (for str, int, float) or pandas scalar
443 (for Timestamp/Timedelta/Interval/Period)
444
445 Returns
446 ----------
447 list
448
449 Examples
450 ----------
451 >>> arr = pd.array(["a", "b", "c"])
452 >>> arr.tolist()
453 ['a', 'b', 'c']
454 """
455 if self.ndim > 1:
456 return [x.tolist() for x in self]
457 return list(self.to_numpy())
458
459 def _formatter(self, boxed: bool = False):
460 formatter = partial(
461 printing.pprint_thing,
462 escape_chars=("\t", "\r", "\n"),
463 quote_strings=not boxed,
464 )
465 return formatter
466
467 def _str_map(
468 self,
469 f,
470 na_value=lib.no_default,
471 dtype: Dtype | None = None,
472 convert: bool = True,
473 ):
474 if self.dtype.na_value is np.nan:
475 return self._str_map_nan_semantics(f, na_value=na_value, dtype=dtype)
476
477 from pandas.arrays import BooleanArray
478
479 if dtype is None:
480 dtype = self.dtype
481 if na_value is lib.no_default:
482 na_value = self.dtype.na_value
483
484 mask = isna(self)
485 arr = np.asarray(self)
486
487 if is_integer_dtype(dtype) or is_bool_dtype(dtype):
488 constructor: type[IntegerArray | BooleanArray]
489 if is_integer_dtype(dtype):
490 constructor = IntegerArray
491 else:
492 constructor = BooleanArray
493
494 na_value_is_na = isna(na_value)
495 if na_value_is_na:
496 na_value = 1
497 elif dtype == np.dtype("bool"):
498 # GH#55736
499 na_value = bool(na_value)
500 result = lib.map_infer_mask(
501 arr,
502 f,
503 mask.view("uint8"),
504 convert=False,
505 na_value=na_value,
506 # error: Argument 1 to "dtype" has incompatible type
507 # "Union[ExtensionDtype, str, dtype[Any], Type[object]]"; expected
508 # "Type[object]"
509 dtype=np.dtype(cast("type", dtype)),
510 )
511
512 if not na_value_is_na:
513 mask[:] = False
514
515 return constructor(result, mask)
516
517 else:
518 return self._str_map_str_or_object(dtype, na_value, arr, f, mask)
519
520 def _str_map_str_or_object(
521 self,
522 dtype,
523 na_value,
524 arr: np.ndarray,
525 f,
526 mask: npt.NDArray[np.bool_],
527 ):
528 # _str_map helper for case where dtype is either string dtype or object
529 if is_string_dtype(dtype) and not is_object_dtype(dtype):
530 # i.e. StringDtype
531 result = lib.map_infer_mask(
532 arr, f, mask.view("uint8"), convert=False, na_value=na_value
533 )
534 if self.dtype.storage == "pyarrow":
535 import pyarrow as pa
536
537 # TODO: shouldn't this already be caught my passed mask?
538 # it isn't in test_extract_expand_capture_groups_index
539 # mask = mask | np.array(
540 # [x is libmissing.NA for x in result], dtype=bool
541 # )
542
543 result = pa.array(
544 result, mask=mask, type=pa.large_string(), from_pandas=True
545 )
546 # error: "BaseStringArray" has no attribute "_from_pyarrow_array"
547 return self._from_pyarrow_array(result) # type: ignore[attr-defined]
548 else:
549 # StringArray
550 # error: Too many arguments for "BaseStringArray"
551 return type(self)(result, dtype=self.dtype) # type: ignore[call-arg]
552
553 else:
554 # This is when the result type is object. We reach this when
555 # -> We know the result type is truly object (e.g. .encode returns bytes
556 # or .findall returns a list).
557 # -> We don't know the result type. E.g. `.get` can return anything.
558 return lib.map_infer_mask(arr, f, mask.view("uint8"))
559
560 def _str_map_nan_semantics(
561 self, f, na_value=lib.no_default, dtype: Dtype | None = None
562 ):
563 if dtype is None:
564 dtype = self.dtype
565 if na_value is lib.no_default:
566 if is_bool_dtype(dtype):
567 # NaN propagates as False
568 na_value = False
569 else:
570 na_value = self.dtype.na_value
571
572 mask = isna(self)
573 arr = np.asarray(self)
574
575 if is_integer_dtype(dtype) or is_bool_dtype(dtype):
576 na_value_is_na = isna(na_value)
577 if na_value_is_na:
578 if is_integer_dtype(dtype):
579 na_value = 0
580 else:
581 # NaN propagates as False
582 na_value = False
583
584 result = lib.map_infer_mask(
585 arr,
586 f,
587 mask.view("uint8"),
588 convert=False,
589 na_value=na_value,
590 dtype=np.dtype(cast("type", dtype)),
591 )
592 if na_value_is_na and is_integer_dtype(dtype) and mask.any():
593 # TODO: we could alternatively do this check before map_infer_mask
594 # and adjust the dtype/na_value we pass there. Which is more
595 # performant?
596 result = result.astype("float64")
597 result[mask] = np.nan
598
599 return result
600
601 else:
602 return self._str_map_str_or_object(dtype, na_value, arr, f, mask)
603
604 @overload
605 def view(self, dtype: None = ...) -> Self: ...
606
607 @overload
608 def view(self, dtype: Dtype | None = ...) -> ArrayLike: ...
609
610 def view(self, dtype: Dtype | None = None) -> ArrayLike:
611 if dtype is not None:
612 raise TypeError("Cannot change data-type for string array.")
613 return super().view()
614
615
616@set_module("pandas.arrays")
617# error: Definition of "_concat_same_type" in base class "NDArrayBacked" is
618# incompatible with definition in base class "ExtensionArray"
619class StringArray(BaseStringArray, NumpyExtensionArray): # type: ignore[misc]
620 """
621 Extension array for string data.
622
623 .. warning::
624
625 StringArray is considered experimental. The implementation and
626 parts of the API may change without warning.
627
628 Parameters
629 ----------
630 values : array-like
631 The array of data.
632
633 .. warning::
634
635 Currently, this expects an object-dtype ndarray
636 where the elements are Python strings
637 or nan-likes (``None``, ``np.nan``, ``NA``).
638 This may change without warning in the future. Use
639 :meth:`pandas.array` with ``dtype="string"`` for a stable way of
640 creating a `StringArray` from any sequence.
641
642 StringArray accepts array-likes containing
643 nan-likes(``None``, ``np.nan``) for the ``values`` parameter
644 in addition to strings and :attr:`pandas.NA`
645
646 dtype : StringDtype
647 Dtype for the array.
648 copy : bool, default False
649 Whether to copy the array of data.
650
651 Attributes
652 ----------
653 None
654
655 Methods
656 -------
657 None
658
659 See Also
660 --------
661 :func:`array`
662 The recommended function for creating a StringArray.
663 Series.str
664 The string methods are available on Series backed by
665 a StringArray.
666
667 Notes
668 -----
669 StringArray returns a BooleanArray for comparison methods.
670
671 Examples
672 --------
673 >>> pd.array(["This is", "some text", None, "data."], dtype="string")
674 <ArrowStringArray>
675 ['This is', 'some text', <NA>, 'data.']
676 Length: 4, dtype: string
677
678 Unlike arrays instantiated with ``dtype="object"``, ``StringArray``
679 will convert the values to strings.
680
681 >>> pd.array(["1", 1], dtype="object")
682 <NumpyExtensionArray>
683 ['1', 1]
684 Length: 2, dtype: object
685 >>> pd.array(["1", 1], dtype="string")
686 <ArrowStringArray>
687 ['1', '1']
688 Length: 2, dtype: string
689
690 However, instantiating StringArrays directly with non-strings will raise an error.
691
692 For comparison methods, `StringArray` returns a :class:`pandas.BooleanArray`:
693
694 >>> pd.array(["a", None, "c"], dtype="string[python]") == "a"
695 <BooleanArray>
696 [True, <NA>, False]
697 Length: 3, dtype: boolean
698 """
699
700 # undo the NumpyExtensionArray hack
701 _typ = "extension"
702
703 def __init__(
704 self, values, *, dtype: StringDtype | None = None, copy: bool = False
705 ) -> None:
706 if dtype is None:
707 dtype = StringDtype()
708 values = extract_array(values)
709
710 super().__init__(values, copy=copy)
711 if not isinstance(values, type(self)):
712 self._validate(dtype)
713 NDArrayBacked.__init__(
714 self,
715 self._ndarray,
716 dtype,
717 )
718
719 def _validate(self, dtype: StringDtype) -> None:
720 """Validate that we only store NA or strings."""
721
722 if dtype._na_value is libmissing.NA:
723 if len(self._ndarray) and not lib.is_string_array(
724 self._ndarray, skipna=True
725 ):
726 raise ValueError(
727 "StringArray requires a sequence of strings or pandas.NA"
728 )
729 if self._ndarray.dtype != "object":
730 raise ValueError(
731 "StringArray requires a sequence of strings or pandas.NA. Got "
732 f"'{self._ndarray.dtype}' dtype instead."
733 )
734 # Check to see if need to convert Na values to pd.NA
735 if self._ndarray.ndim > 2:
736 # Ravel if ndims > 2 b/c no cythonized version available
737 lib.convert_nans_to_NA(self._ndarray.ravel("K"))
738 else:
739 lib.convert_nans_to_NA(self._ndarray)
740 else:
741 # Validate that we only store NaN or strings.
742 if len(self._ndarray) and not lib.is_string_array(
743 self._ndarray, skipna=True
744 ):
745 raise ValueError("StringArray requires a sequence of strings or NaN")
746 if self._ndarray.dtype != "object":
747 raise ValueError(
748 "StringArray requires a sequence of strings "
749 "or NaN. Got '{self._ndarray.dtype}' dtype instead."
750 )
751 # TODO validate or force NA/None to NaN
752
753 def _validate_scalar(self, value):
754 # used by NDArrayBackedExtensionIndex.insert
755 if isna(value):
756 return self.dtype.na_value
757 elif not isinstance(value, str):
758 raise TypeError(
759 f"Invalid value '{value}' for dtype '{self.dtype}'. Value should be a "
760 f"string or missing value, got '{type(value).__name__}' instead."
761 )
762 return value
763
764 @classmethod
765 def _from_sequence(
766 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
767 ) -> Self:
768 if dtype and not (isinstance(dtype, str) and dtype == "string"):
769 dtype = pandas_dtype(dtype)
770 assert isinstance(dtype, StringDtype) and dtype.storage == "python"
771 elif using_string_dtype():
772 dtype = StringDtype(storage="python", na_value=np.nan)
773 else:
774 dtype = StringDtype(storage="python")
775
776 from pandas.core.arrays.masked import BaseMaskedArray
777
778 na_value = dtype.na_value
779 if isinstance(scalars, BaseMaskedArray):
780 # avoid costly conversion to object dtype
781 na_values = scalars._mask
782 result = scalars._data
783 result = lib.ensure_string_array(
784 result, copy=copy, convert_na_value=False, skipna=False
785 )
786 result[na_values] = na_value
787
788 else:
789 if lib.is_pyarrow_array(scalars):
790 # pyarrow array; we cannot rely on the "to_numpy" check in
791 # ensure_string_array because calling scalars.to_numpy would set
792 # zero_copy_only to True which caused problems see GH#52076
793 scalars = np.array(scalars)
794 # convert non-na-likes to str, and nan-likes to StringDtype().na_value
795 result = lib.ensure_string_array(scalars, na_value=na_value, copy=copy)
796
797 # Manually creating new array avoids the validation step in the __init__, so is
798 # faster. Refactor need for validation?
799 new_string_array = cls.__new__(cls)
800 NDArrayBacked.__init__(new_string_array, result, dtype)
801
802 return new_string_array
803
804 @classmethod
805 def _from_sequence_of_strings(
806 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
807 ) -> Self:
808 return cls._from_sequence(strings, dtype=dtype, copy=copy)
809
810 def _cast_pointwise_result(self, values) -> ArrayLike:
811 result = super()._cast_pointwise_result(values)
812 if isinstance(result.dtype, StringDtype):
813 # Ensure we retain our same na_value/storage
814 result = result.astype(self.dtype)
815 return result
816
817 @classmethod
818 def _empty(cls, shape, dtype) -> StringArray:
819 values = np.empty(shape, dtype=object)
820 values[:] = dtype.na_value
821 return cls(values, dtype=dtype).astype(dtype, copy=False)
822
823 def __arrow_array__(self, type=None):
824 """
825 Convert myself into a pyarrow Array.
826 """
827 import pyarrow as pa
828
829 if type is None:
830 type = pa.string()
831
832 values = self._ndarray.copy()
833 values[self.isna()] = None
834 return pa.array(values, type=type)
835
836 def isna(self) -> np.ndarray:
837 return libmissing.isna_string(self._ndarray)
838
839 def _values_for_factorize(self) -> tuple[np.ndarray, libmissing.NAType | float]: # type: ignore[override]
840 arr = self._ndarray
841
842 return arr, self.dtype.na_value
843
844 def _validate_setitem_value(self, value):
845 """Validate / convert value to be StringArray compatible."""
846 if lib.is_scalar(value):
847 if isna(value):
848 value = self.dtype.na_value
849 elif not isinstance(value, str):
850 raise TypeError(
851 f"Invalid value '{value}' for dtype '{self.dtype}'. Value should "
852 f"be a string or missing value, got '{type(value).__name__}' "
853 "instead."
854 )
855 else:
856 value = extract_array(value, extract_numpy=True)
857 if not is_array_like_deprecate_non_pandas(value):
858 value = np.asarray(value, dtype=object)
859 elif isinstance(value.dtype, type(self.dtype)):
860 return value
861 else:
862 # cast categories and friends to arrays to see if values are
863 # compatible, compatibility with arrow backed strings
864 value = np.asarray(value)
865 if len(value) and not lib.is_string_array(value, skipna=True):
866 raise TypeError(
867 "Invalid value for dtype 'str'. Value should be a "
868 "string or missing value (or array of those)."
869 )
870 return value
871
872 def __setitem__(self, key, value) -> None:
873 if self._readonly:
874 raise ValueError("Cannot modify read-only array")
875
876 value = self._validate_setitem_value(value)
877
878 key = check_array_indexer(self, key)
879 scalar_key = lib.is_scalar(key)
880 scalar_value = lib.is_scalar(value)
881 if scalar_key and not scalar_value:
882 raise ValueError("setting an array element with a sequence.")
883
884 if not scalar_value:
885 if value.dtype == self.dtype:
886 value = value._ndarray
887 else:
888 value = np.asarray(value)
889 mask = isna(value)
890 if mask.any():
891 value = value.copy()
892 value[isna(value)] = self.dtype.na_value
893
894 super().__setitem__(key, value)
895
896 def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
897 # the super() method NDArrayBackedExtensionArray._putmask uses
898 # np.putmask which doesn't properly handle None/pd.NA, so using the
899 # base class implementation that uses __setitem__
900 ExtensionArray._putmask(self, mask, value)
901
902 def _where(self, mask: npt.NDArray[np.bool_], value) -> Self:
903 # the super() method NDArrayBackedExtensionArray._where uses
904 # np.putmask which doesn't properly handle None/pd.NA, so using the
905 # base class implementation that uses __setitem__
906 return ExtensionArray._where(self, mask, value)
907
908 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
909 if isinstance(values, BaseStringArray) or (
910 isinstance(values, ExtensionArray) and is_string_dtype(values.dtype)
911 ):
912 values = values.astype(self.dtype, copy=False)
913 else:
914 if not lib.is_string_array(np.asarray(values), skipna=True):
915 values = np.array(
916 [val for val in values if isinstance(val, str) or isna(val)],
917 dtype=object,
918 )
919 if not len(values):
920 return np.zeros(self.shape, dtype=bool)
921
922 values = self._from_sequence(values, dtype=self.dtype)
923
924 return isin(np.asarray(self), np.asarray(values))
925
926 def astype(self, dtype, copy: bool = True):
927 dtype = pandas_dtype(dtype)
928
929 if dtype == self.dtype:
930 if copy:
931 return self.copy()
932 return self
933
934 elif isinstance(dtype, IntegerDtype):
935 arr = self._ndarray.copy()
936 mask = self.isna()
937 arr[mask] = 0
938 values = arr.astype(dtype.numpy_dtype)
939 return IntegerArray(values, mask, copy=False)
940 elif isinstance(dtype, FloatingDtype):
941 arr_ea = self.copy()
942 mask = self.isna()
943 arr_ea[mask] = "0"
944 values = arr_ea.astype(dtype.numpy_dtype)
945 return FloatingArray(values, mask, copy=False)
946 elif isinstance(dtype, ExtensionDtype):
947 # Skip the NumpyExtensionArray.astype method
948 return ExtensionArray.astype(self, dtype, copy)
949 elif np.issubdtype(dtype, np.floating):
950 arr = self._ndarray.copy()
951 mask = self.isna()
952 arr[mask] = 0
953 values = arr.astype(dtype)
954 values[mask] = np.nan
955 return values
956
957 return super().astype(dtype, copy)
958
959 def _reduce(
960 self,
961 name: str,
962 *,
963 skipna: bool = True,
964 keepdims: bool = False,
965 axis: AxisInt | None = 0,
966 **kwargs,
967 ):
968 if self.dtype.na_value is np.nan and name in ["any", "all"]:
969 if name == "any":
970 return nanops.nanany(self._ndarray, skipna=skipna)
971 else:
972 return nanops.nanall(self._ndarray, skipna=skipna)
973 elif name == "count":
974 return super().count()
975 elif name in ["min", "max", "argmin", "argmax", "sum"]:
976 result = getattr(self, name)(skipna=skipna, axis=axis, **kwargs)
977 if keepdims:
978 return self._from_sequence([result], dtype=self.dtype)
979 return result
980
981 raise TypeError(f"Cannot perform reduction '{name}' with string dtype")
982
983 def _accumulate(self, name: str, *, skipna: bool = True, **kwargs) -> StringArray:
984 """
985 Return an ExtensionArray performing an accumulation operation.
986
987 The underlying data type might change.
988
989 Parameters
990 ----------
991 name : str
992 Name of the function, supported values are:
993 - cummin
994 - cummax
995 - cumsum
996 - cumprod
997 skipna : bool, default True
998 If True, skip NA values.
999 **kwargs
1000 Additional keyword arguments passed to the accumulation function.
1001 Currently, there is no supported kwarg.
1002
1003 Returns
1004 -------
1005 array
1006
1007 Raises
1008 ------
1009 NotImplementedError : subclass does not define accumulations
1010 """
1011 if name == "cumprod":
1012 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
1013 raise TypeError(msg)
1014
1015 # We may need to strip out trailing NA values
1016 tail: np.ndarray | None = None
1017 na_mask: np.ndarray | None = None
1018 ndarray = self._ndarray
1019 np_func = {
1020 "cumsum": np.cumsum,
1021 "cummin": np.minimum.accumulate,
1022 "cummax": np.maximum.accumulate,
1023 }[name]
1024
1025 if self._hasna:
1026 na_mask = cast("npt.NDArray[np.bool_]", isna(ndarray))
1027 if np.all(na_mask):
1028 return type(self)(ndarray, dtype=self.dtype)
1029 if skipna:
1030 if name == "cumsum":
1031 ndarray = np.where(na_mask, "", ndarray)
1032 else:
1033 # We can retain the running min/max by forward/backward filling.
1034 ndarray = ndarray.copy()
1035 missing.pad_or_backfill_inplace(
1036 ndarray,
1037 method="pad",
1038 axis=0,
1039 )
1040 missing.pad_or_backfill_inplace(
1041 ndarray,
1042 method="backfill",
1043 axis=0,
1044 )
1045 else:
1046 # When not skipping NA values, the result should be null from
1047 # the first NA value onward.
1048 idx = np.argmax(na_mask)
1049 tail = np.empty(len(ndarray) - idx, dtype="object")
1050 tail[:] = self.dtype.na_value
1051 ndarray = ndarray[:idx]
1052
1053 # mypy: Cannot call function of unknown type
1054 np_result = np_func(ndarray) # type: ignore[operator]
1055
1056 if tail is not None:
1057 np_result = np.hstack((np_result, tail))
1058 elif na_mask is not None:
1059 # Argument 2 to "where" has incompatible type "NAType | float"
1060 np_result = np.where(na_mask, self.dtype.na_value, np_result) # type: ignore[arg-type]
1061
1062 result = type(self)(np_result, dtype=self.dtype)
1063 return result
1064
1065 def _wrap_reduction_result(self, axis: AxisInt | None, result) -> Any:
1066 if self.dtype.na_value is np.nan and result is libmissing.NA:
1067 # the masked_reductions use pd.NA -> convert to np.nan
1068 return np.nan
1069 return super()._wrap_reduction_result(axis, result)
1070
1071 def any(
1072 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
1073 ) -> bool:
1074 nv.validate_any((), kwargs)
1075 return self._reduce("any", axis=axis, skipna=skipna, **kwargs)
1076
1077 def all(
1078 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
1079 ) -> bool:
1080 nv.validate_all((), kwargs)
1081 return self._reduce("all", axis=axis, skipna=skipna, **kwargs)
1082
1083 def min(
1084 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
1085 ) -> Scalar:
1086 nv.validate_min((), kwargs)
1087 result = masked_reductions.min(
1088 values=self.to_numpy(), mask=self.isna(), skipna=skipna
1089 )
1090 return self._wrap_reduction_result(axis, result)
1091
1092 def max(
1093 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
1094 ) -> Scalar:
1095 nv.validate_max((), kwargs)
1096 result = masked_reductions.max(
1097 values=self.to_numpy(), mask=self.isna(), skipna=skipna
1098 )
1099 return self._wrap_reduction_result(axis, result)
1100
1101 def sum(
1102 self,
1103 *,
1104 axis: AxisInt | None = None,
1105 skipna: bool = True,
1106 min_count: int = 0,
1107 **kwargs,
1108 ) -> Scalar:
1109 nv.validate_sum((), kwargs)
1110 result = masked_reductions.sum(
1111 values=self._ndarray,
1112 mask=self.isna(),
1113 skipna=skipna,
1114 min_count=min_count,
1115 initial="",
1116 )
1117 return self._wrap_reduction_result(axis, result)
1118
1119 # Without these, NumpyExtensionArray's methods bypass _reduce's gate, GH#68389
1120 def prod(self, *, skipna: bool = True, **kwargs) -> Scalar:
1121 return self._reduce("prod", skipna=skipna, **kwargs)
1122
1123 def mean(self, *, skipna: bool = True, **kwargs) -> Scalar:
1124 return self._reduce("mean", skipna=skipna, **kwargs)
1125
1126 def median(self, *, skipna: bool = True, **kwargs) -> Scalar:
1127 return self._reduce("median", skipna=skipna, **kwargs)
1128
1129 def std(self, *, skipna: bool = True, **kwargs) -> Scalar:
1130 return self._reduce("std", skipna=skipna, **kwargs)
1131
1132 def var(self, *, skipna: bool = True, **kwargs) -> Scalar:
1133 return self._reduce("var", skipna=skipna, **kwargs)
1134
1135 def sem(self, *, skipna: bool = True, **kwargs) -> Scalar:
1136 return self._reduce("sem", skipna=skipna, **kwargs)
1137
1138 def skew(self, *, skipna: bool = True, **kwargs) -> Scalar:
1139 return self._reduce("skew", skipna=skipna, **kwargs)
1140
1141 def kurt(self, *, skipna: bool = True, **kwargs) -> Scalar:
1142 return self._reduce("kurt", skipna=skipna, **kwargs)
1143
1144 def value_counts(self, dropna: bool = True) -> Series:
1145 result = super().value_counts(dropna=dropna)
1146
1147 if self.dtype.na_value is libmissing.NA:
1148 result = result.astype("Int64")
1149 return result
1150
1151 def memory_usage(self, deep: bool = False) -> int:
1152 result = self._ndarray.nbytes
1153 if deep and not PYPY:
1154 # deep introspection relies on sys.getsizeof, which always
1155 # raises TypeError on PyPy (GH#46176)
1156 return result + lib.memory_usage_of_objects(self._ndarray)
1157 return result
1158
1159 def searchsorted(
1160 self,
1161 value: NumpyValueArrayLike | ExtensionArray,
1162 side: Literal["left", "right"] = "left",
1163 sorter: NumpySorter | None = None,
1164 ) -> npt.NDArray[np.intp] | np.intp:
1165 """
1166 Find indices where elements should be inserted to maintain order.
1167
1168 Find the indices into a sorted array `self` (a) such that, if the
1169 corresponding elements in `value` were inserted before the indices,
1170 the order of `self` would be preserved.
1171
1172 Assuming that `self` is sorted:
1173
1174 ====== ================================
1175 `side` returned index `i` satisfies
1176 ====== ================================
1177 left ``self[i-1] < value <= self[i]``
1178 right ``self[i-1] <= value < self[i]``
1179 ====== ================================
1180
1181 Parameters
1182 ----------
1183 value : array-like, list or scalar
1184 Value(s) to insert into `self`.
1185 side : {'left', 'right'}, optional
1186 If 'left', the index of the first suitable location found is given.
1187 If 'right', return the last such index. If there is no suitable
1188 index, return either 0 or N (where N is the length of `self`).
1189 sorter : 1-D array-like, optional
1190 Optional array of integer indices that sort array a into ascending
1191 order. They are typically the result of argsort.
1192
1193 Returns
1194 -------
1195 array of ints or int
1196 If value is array-like, array of insertion points.
1197 If value is scalar, a single integer.
1198
1199 See Also
1200 --------
1201 numpy.searchsorted : Similar method from NumPy.
1202
1203 Examples
1204 --------
1205 >>> arr = pd.array([1, 2, 3, 5])
1206 >>> arr.searchsorted([4])
1207 array([3])
1208 """
1209
1210 # GH#65837: avoid O(n) scan; NA confined to array ends in sorted data.
1211 # When sorter is given, the sorted order is ndarray[sorter], so check
1212 # the first/last positions via sorter instead of raw ndarray positions.
1213 ndarray = self._ndarray
1214 if len(ndarray):
1215 if sorter is None:
1216 has_na = libmissing.checknull(ndarray[0]) or libmissing.checknull(
1217 ndarray[-1]
1218 )
1219 else:
1220 has_na = libmissing.checknull(
1221 ndarray[sorter[0]]
1222 ) or libmissing.checknull(ndarray[sorter[-1]])
1223 else:
1224 has_na = False
1225 if has_na:
1226 raise ValueError(
1227 "searchsorted requires array to be sorted, which is impossible "
1228 "with NAs present."
1229 )
1230 return super().searchsorted(value=value, side=side, sorter=sorter)
1231
1232 def _cmp_method(self, other, op):
1233 from pandas.arrays import (
1234 ArrowExtensionArray,
1235 BooleanArray,
1236 )
1237
1238 ops.raise_if_2d(other)
1239
1240 if (
1241 isinstance(other, BaseStringArray)
1242 and self.dtype.na_value is not libmissing.NA
1243 and other.dtype.na_value is libmissing.NA
1244 ):
1245 # NA has priority of NaN semantics
1246 return op(self.astype(other.dtype, copy=False), other)
1247
1248 if isinstance(other, ArrowExtensionArray):
1249 if isinstance(other, BaseStringArray):
1250 # pyarrow storage has priority over python storage
1251 # (except if we have NA semantics and other not)
1252 if not (
1253 self.dtype.na_value is libmissing.NA
1254 and other.dtype.na_value is not libmissing.NA
1255 ):
1256 return NotImplemented
1257 else:
1258 return NotImplemented
1259
1260 if isinstance(other, StringArray):
1261 other = other._ndarray
1262
1263 mask = isna(self) | isna(other)
1264 valid = ~mask
1265
1266 if lib.is_list_like(other):
1267 if not isinstance(
1268 other, (list, ExtensionArray, np.ndarray)
1269 ) and not ops.has_castable_attr(other):
1270 warnings.warn(
1271 f"Operation with {type(other).__name__} is deprecated. "
1272 "In a future version these will be treated as scalar-like. "
1273 "To retain the old behavior, explicitly wrap in a Series "
1274 "instead.",
1275 Pandas4Warning,
1276 stacklevel=find_stack_level(),
1277 )
1278 if len(other) != len(self):
1279 raise ValueError(
1280 f"Lengths of operands do not match: {len(self)} != {len(other)}"
1281 )
1282
1283 # for array-likes, first filter out NAs before converting to numpy
1284 if not is_array_like_deprecate_non_pandas(other):
1285 other = np.asarray(other)
1286 other = other[valid]
1287
1288 other_dtype = getattr(other, "dtype", None)
1289 if op.__name__.strip("_") in ["mul", "rmul"] and (
1290 lib.is_bool(other) or lib.is_np_dtype(other_dtype, "b")
1291 ):
1292 # GH#62595
1293 raise TypeError(
1294 "Cannot multiply StringArray by bools. "
1295 "Explicitly cast to integers instead."
1296 )
1297
1298 if op.__name__ in ops.ARITHMETIC_BINOPS:
1299 result = np.empty_like(self._ndarray, dtype="object")
1300 result[mask] = self.dtype.na_value
1301 result[valid] = op(self._ndarray[valid], other)
1302 # GH#40624 is_string_array is False for a length-zero array, but a
1303 # length-zero result should keep the string dtype
1304 if len(result) and not lib.is_string_array(result, skipna=True):
1305 return result
1306 return self._from_backing_data(result)
1307 else:
1308 # logical
1309 result = np.zeros(len(self._ndarray), dtype="bool")
1310 result[valid] = op(self._ndarray[valid], other)
1311 res_arr = BooleanArray(result, mask)
1312 if self.dtype.na_value is np.nan:
1313 if op == operator.ne:
1314 return res_arr.to_numpy(np.bool_, na_value=True)
1315 else:
1316 return res_arr.to_numpy(np.bool_, na_value=False)
1317 return res_arr
1318
1319 _arith_method = _cmp_method
1320
1321 def _str_zfill(self, width: int) -> Self:
1322 return self._str_map(lambda x: x.zfill(width))