rename storage option and add na_value keyword

Author: jorisvandenbossche

pandas/core/arrays/string_.py

@@ -9,7 +9,10 @@
 
 import numpy as np
 
-from pandas._config import get_option
+from pandas._config import (
+    get_option,
+    using_pyarrow_string_dtype,
+)
 
 from pandas._libs import (
     lib,
@@ -83,6 +86,7 @@ class StringDtype(StorageExtensionDtype):
     ----------
     storage : {"python", "pyarrow", "pyarrow_numpy"}, optional
         If not given, the value of ``pd.options.mode.string_storage``.
+    na_value :
 
     Attributes
     ----------
@@ -113,30 +117,49 @@ class StringDtype(StorageExtensionDtype):
     # follows NumPy semantics, which uses nan.
     @property
     def na_value(self) -> libmissing.NAType | float:  # type: ignore[override]
-        if self.storage == "pyarrow_numpy":
-            return np.nan
-        else:
-            return libmissing.NA
+        return self._na_value
 
     _metadata = ("storage",)
 
-    def __init__(self, storage=None) -> None:
-        if storage is None:
-            infer_string = get_option("future.infer_string")
-            if infer_string:
-                storage = "pyarrow_numpy"
+    def __init__(self, storage=None, na_value=None) -> None:
+        if not (
+            na_value is None or (isinstance(na_value, float) and np.isnan(na_value))
+        ):
+            raise ValueError(
+                "'na_value' must be the default value or pd.NA, got {na_value}"
+            )
+
+        # infer defaults
+        if storage is None and na_value is None:
+            if using_pyarrow_string_dtype():
+                storage = "pyarrow"
+                na_value = np.nan
             else:
                 storage = get_option("mode.string_storage")
-        if storage not in {"python", "pyarrow", "pyarrow_numpy"}:
+                na_value = libmissing.NA
+        elif storage is None:
+            # in this case na_value is NaN
+            storage = get_option("mode.string_storage")
+        elif na_value is None:
+            na_value = np.nan if using_pyarrow_string_dtype() else libmissing.NA
+            if na_value is not libmissing.NA and storage == "python":
+                raise NotImplementedError(
+                    "'python' mode for na_value of NaN not yet implemented"
+                )
+
+        if storage == "pyarrow_numpy":
+            # TODO raise a deprecation warning
+            storage = "pyarrow"
+        if storage not in {"python", "pyarrow"}:
             raise ValueError(
-                f"Storage must be 'python', 'pyarrow' or 'pyarrow_numpy'. "
-                f"Got {storage} instead."
+                f"Storage must be 'python' or 'pyarrow'. Got {storage} instead."
             )
-        if storage in ("pyarrow", "pyarrow_numpy") and pa_version_under10p1:
+        if storage == "pyarrow" and pa_version_under10p1:
             raise ImportError(
                 "pyarrow>=10.0.1 is required for PyArrow backed StringArray."
             )
         self.storage = storage
+        self._na_value = na_value
 
     @property
     def type(self) -> type[str]:
@@ -176,11 +199,14 @@ def construct_from_string(cls, string) -> Self:
             )
         if string == "string":
             return cls()
+        elif string == "String":
+            return cls(na_value=np.nan)
         elif string == "string[python]":
-            return cls(storage="python")
+            return cls(storage="python", na_value=np.nan)
         elif string == "string[pyarrow]":
-            return cls(storage="pyarrow")
+            return cls(storage="pyarrow", na_value=np.nan)
         elif string == "string[pyarrow_numpy]":
+            # TODO deprecate
             return cls(storage="pyarrow_numpy")
         else:
             raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
@@ -205,10 +231,10 @@ def construct_array_type(  # type: ignore[override]
 
         if self.storage == "python":
             return StringArray
-        elif self.storage == "pyarrow":
-            return ArrowStringArray
-        else:
+        elif self.storage == "pyarrow" and self._na_value is libmissing.NA:
             return ArrowStringArrayNumpySemantics
+        else:
+            return ArrowStringArray
 
     def __from_arrow__(
         self, array: pyarrow.Array | pyarrow.ChunkedArray
@@ -217,13 +243,16 @@ def __from_arrow__(
         Construct StringArray from pyarrow Array/ChunkedArray.
         """
         if self.storage == "pyarrow":
-            from pandas.core.arrays.string_arrow import ArrowStringArray
+            if self._na_value is libmissing.NA:
+                from pandas.core.arrays.string_arrow import (
+                    ArrowStringArrayNumpySemantics,
+                )
 
-            return ArrowStringArray(array)
-        elif self.storage == "pyarrow_numpy":
-            from pandas.core.arrays.string_arrow import ArrowStringArrayNumpySemantics
+                return ArrowStringArrayNumpySemantics(array)
+            else:
+                from pandas.core.arrays.string_arrow import ArrowStringArray
 
-            return ArrowStringArrayNumpySemantics(array)
+                return ArrowStringArray(array)
         else:
             import pyarrow