[ArrowStringDtype] Make it already a StringDtype subclass (pandas-dev#41312)

Author: jorisvandenbossche

pandas/core/arrays/base.py

@@ -530,7 +530,6 @@ def astype(self, dtype, copy=True):
             NumPy ndarray with 'dtype' for its dtype.
         """
         from pandas.core.arrays.string_ import StringDtype
-        from pandas.core.arrays.string_arrow import ArrowStringDtype
 
         dtype = pandas_dtype(dtype)
         if is_dtype_equal(dtype, self.dtype):
@@ -540,9 +539,8 @@ def astype(self, dtype, copy=True):
                 return self.copy()
 
         # FIXME: Really hard-code here?
-        if isinstance(
-            dtype, (ArrowStringDtype, StringDtype)
-        ):  # allow conversion to StringArrays
+        if isinstance(dtype, StringDtype):
+            # allow conversion to StringArrays
             return dtype.construct_array_type()._from_sequence(self, copy=False)
 
         return np.array(self, dtype=dtype, copy=copy)

pandas/core/arrays/interval.py

@@ -829,7 +829,6 @@ def astype(self, dtype, copy: bool = True):
         """
         from pandas import Index
         from pandas.core.arrays.string_ import StringDtype
-        from pandas.core.arrays.string_arrow import ArrowStringDtype
 
         if dtype is not None:
             dtype = pandas_dtype(dtype)
@@ -852,7 +851,7 @@ def astype(self, dtype, copy: bool = True):
             return self._shallow_copy(new_left, new_right)
         elif is_categorical_dtype(dtype):
             return Categorical(np.asarray(self), dtype=dtype)
-        elif isinstance(dtype, (StringDtype, ArrowStringDtype)):
+        elif isinstance(dtype, StringDtype):
             return dtype.construct_array_type()._from_sequence(self, copy=False)
 
         # TODO: This try/except will be repeated.

pandas/core/arrays/string_arrow.py

@@ -24,7 +24,6 @@
 from pandas.util._decorators import doc
 from pandas.util._validators import validate_fillna_kwargs
 
-from pandas.core.dtypes.base import ExtensionDtype
 from pandas.core.dtypes.common import (
     is_array_like,
     is_bool_dtype,
@@ -42,6 +41,7 @@
 from pandas.core.arrays.base import ExtensionArray
 from pandas.core.arrays.boolean import BooleanDtype
 from pandas.core.arrays.integer import Int64Dtype
+from pandas.core.arrays.string_ import StringDtype
 from pandas.core.indexers import (
     check_array_indexer,
     validate_indices,
@@ -74,7 +74,7 @@
 
 
 @register_extension_dtype
-class ArrowStringDtype(ExtensionDtype):
+class ArrowStringDtype(StringDtype):
     """
     Extension dtype for string data in a ``pyarrow.ChunkedArray``.
 
@@ -110,7 +110,7 @@ def type(self) -> type[str]:
         return str
 
     @classmethod
-    def construct_array_type(cls) -> type_t[ArrowStringArray]:
+    def construct_array_type(cls) -> type_t[ArrowStringArray]:  # type: ignore[override]
         """
         Return the array type associated with this dtype.
 
@@ -126,7 +126,9 @@ def __hash__(self) -> int:
     def __repr__(self) -> str:
         return "ArrowStringDtype"
 
-    def __from_arrow__(self, array: pa.Array | pa.ChunkedArray) -> ArrowStringArray:
+    def __from_arrow__(  # type: ignore[override]
+        self, array: pa.Array | pa.ChunkedArray
+    ) -> ArrowStringArray:
         """
         Construct StringArray from pyarrow Array/ChunkedArray.
         """

pandas/core/strings/accessor.py

@@ -155,11 +155,10 @@ class StringMethods(NoNewAttributesMixin):
 
     def __init__(self, data):
         from pandas.core.arrays.string_ import StringDtype
-        from pandas.core.arrays.string_arrow import ArrowStringDtype
 
         self._inferred_dtype = self._validate(data)
         self._is_categorical = is_categorical_dtype(data.dtype)
-        self._is_string = isinstance(data.dtype, (StringDtype, ArrowStringDtype))
+        self._is_string = isinstance(data.dtype, StringDtype)
         self._data = data
 
         self._index = self._name = None
@@ -3028,9 +3027,8 @@ def _result_dtype(arr):
     # ideally we just pass `dtype=arr.dtype` unconditionally, but this fails
     # when the list of values is empty.
     from pandas.core.arrays.string_ import StringDtype
-    from pandas.core.arrays.string_arrow import ArrowStringDtype
 
-    if isinstance(arr.dtype, (StringDtype, ArrowStringDtype)):
+    if isinstance(arr.dtype, StringDtype):
         return arr.dtype.name
     else:
         return object

pandas/tests/extension/json/array.py

@@ -40,7 +40,6 @@
     ExtensionDtype,
 )
 from pandas.api.types import is_bool_dtype
-from pandas.core.arrays.string_arrow import ArrowStringDtype
 
 
 class JSONDtype(ExtensionDtype):
@@ -196,7 +195,7 @@ def astype(self, dtype, copy=True):
             if copy:
                 return self.copy()
             return self
-        elif isinstance(dtype, (StringDtype, ArrowStringDtype)):
+        elif isinstance(dtype, StringDtype):
             value = self.astype(str)  # numpy doesn'y like nested dicts
             return dtype.construct_array_type()._from_sequence(value, copy=False)