REF: de-duplicate ensure_np_dtype (#54258)

* REF: de-duplicate ensure_np_dtype

* REF: simplify as_array

* mypy fixup

* Missing check in CoW case

Author: jbrockmendel

pandas/core/internals/array_manager.py

@@ -35,11 +35,7 @@
     is_object_dtype,
     is_timedelta64_ns_dtype,
 )
-from pandas.core.dtypes.dtypes import (
-    ExtensionDtype,
-    NumpyEADtype,
-    SparseDtype,
-)
+from pandas.core.dtypes.dtypes import ExtensionDtype
 from pandas.core.dtypes.generic import (
     ABCDataFrame,
     ABCSeries,
@@ -75,6 +71,7 @@
 from pandas.core.internals.base import (
     DataManager,
     SingleDataManager,
+    ensure_np_dtype,
     interleaved_dtype,
 )
 from pandas.core.internals.blocks import (
@@ -1021,14 +1018,7 @@ def as_array(
         if not dtype:
             dtype = interleaved_dtype([arr.dtype for arr in self.arrays])
 
-        if isinstance(dtype, SparseDtype):
-            dtype = dtype.subtype
-        elif isinstance(dtype, NumpyEADtype):
-            dtype = dtype.numpy_dtype
-        elif isinstance(dtype, ExtensionDtype):
-            dtype = np.dtype("object")
-        elif dtype == np.dtype(str):
-            dtype = np.dtype("object")
+        dtype = ensure_np_dtype(dtype)
 
         result = np.empty(self.shape_proper, dtype=dtype)
 

pandas/core/internals/base.py

@@ -8,6 +8,7 @@
     TYPE_CHECKING,
     Any,
     Literal,
+    cast,
     final,
 )
 
@@ -26,6 +27,10 @@
     find_common_type,
     np_can_hold_element,
 )
+from pandas.core.dtypes.dtypes import (
+    ExtensionDtype,
+    SparseDtype,
+)
 
 from pandas.core.base import PandasObject
 from pandas.core.construction import extract_array
@@ -356,3 +361,16 @@ def interleaved_dtype(dtypes: list[DtypeObj]) -> DtypeObj | None:
         return None
 
     return find_common_type(dtypes)
+
+
+def ensure_np_dtype(dtype: DtypeObj) -> np.dtype:
+    # TODO: https://github.com/pandas-dev/pandas/issues/22791
+    # Give EAs some input on what happens here. Sparse needs this.
+    if isinstance(dtype, SparseDtype):
+        dtype = dtype.subtype
+        dtype = cast(np.dtype, dtype)
+    elif isinstance(dtype, ExtensionDtype):
+        dtype = np.dtype("object")
+    elif dtype == np.dtype(str):
+        dtype = np.dtype("object")
+    return dtype

pandas/core/internals/managers.py

@@ -65,6 +65,7 @@
 from pandas.core.internals.base import (
     DataManager,
     SingleDataManager,
+    ensure_np_dtype,
     interleaved_dtype,
 )
 from pandas.core.internals.blocks import (
@@ -1623,16 +1624,12 @@ def as_array(
                 arr = blk.values.to_numpy(  # type: ignore[union-attr]
                     dtype=dtype,
                     na_value=na_value,
+                    copy=copy,
                 ).reshape(blk.shape)
             else:
-                arr = np.asarray(blk.get_values())
-                if dtype:
-                    arr = arr.astype(dtype, copy=copy)
-                    copy = False
+                arr = np.array(blk.values, dtype=dtype, copy=copy)
 
-            if copy:
-                arr = arr.copy()
-            elif using_copy_on_write():
+            if using_copy_on_write() and not copy:
                 arr = arr.view()
                 arr.flags.writeable = False
         else:
@@ -1666,16 +1663,9 @@ def _interleave(
                 [blk.dtype for blk in self.blocks]
             )
 
-        # TODO: https://github.com/pandas-dev/pandas/issues/22791
-        # Give EAs some input on what happens here. Sparse needs this.
-        if isinstance(dtype, SparseDtype):
-            dtype = dtype.subtype
-            dtype = cast(np.dtype, dtype)
-        elif isinstance(dtype, ExtensionDtype):
-            dtype = np.dtype("object")
-        elif dtype == np.dtype(str):
-            dtype = np.dtype("object")
-
+        # error: Argument 1 to "ensure_np_dtype" has incompatible type
+        # "Optional[dtype[Any]]"; expected "Union[dtype[Any], ExtensionDtype]"
+        dtype = ensure_np_dtype(dtype)  # type: ignore[arg-type]
         result = np.empty(self.shape, dtype=dtype)
 
         itemmask = np.zeros(self.shape[0])