ENH: EADtype._find_compatible_dtype

pandas/core/arrays/arrow/dtype.py

@@ -8,10 +8,14 @@
 )
 from decimal import Decimal
 import re
-from typing import TYPE_CHECKING
+from typing import (
+ TYPE_CHECKING,
+ Any,
+)
 
 import numpy as np
 
+from pandas._libs import missing as libmissing
 from pandas._libs.tslibs import (
  Timedelta,
  Timestamp,
@@ -23,6 +27,7 @@
  StorageExtensionDtype,
  register_extension_dtype,
 )
+from pandas.core.dtypes.cast import maybe_promote
 from pandas.core.dtypes.dtypes import CategoricalDtypeType
 
 if not pa_version_under7p0:
@@ -321,3 +326,27 @@ def __from_arrow__(self, array: pa.Array | pa.ChunkedArray):
  array_class = self.construct_array_type()
  arr = array.cast(self.pyarrow_dtype, safe=True)
  return array_class(arr)
+
+ def _maybe_promote(self, item: Any) -> tuple[DtypeObj, Any]:
+ if isinstance(item, pa.Scalar):
+ if not item.is_valid:
+ # TODO: ask joris for help making these checks more robust
+ if item.type == self.pyarrow_dtype:
+ return self, item.as_py()
+ if item.type.to_pandas_dtype() == np.int64 and self.kind == "i":
+ # FIXME: kludge
+ return self, item.as_py()
+
+ item = item.as_py()
+
+ elif item is None or item is libmissing.NA:
+ # TODO: np.nan? use is_valid_na_for_dtype
+ return self, item
+
+ dtype, item = maybe_promote(self.numpy_dtype, item)
+
+ if dtype == self.numpy_dtype:
+ return self, item
+
+ # TODO: implement from_numpy_dtype analogous to MaskedDtype.from_numpy_dtype
+ return np.dtype(object), item

pandas/core/dtypes/base.py

@@ -391,6 +391,9 @@ def _can_hold_na(self) -> bool:
  """
  return True
 
+ def _maybe_promote(self, item: Any) -> tuple[DtypeObj, Any]:
+ return np.dtype(object), item
+
 
 class StorageExtensionDtype(ExtensionDtype):
  """ExtensionDtype that may be backed by more than one implementation."""

pandas/core/dtypes/cast.py

@@ -46,7 +46,6 @@
  ensure_int16,
  ensure_int32,
  ensure_int64,
- ensure_object,
  ensure_str,
  is_bool,
  is_complex,
@@ -539,13 +538,13 @@ def ensure_dtype_can_hold_na(dtype: DtypeObj) -> DtypeObj:
 }
 
 
-def maybe_promote(dtype: np.dtype, fill_value=np.nan):
+def maybe_promote(dtype: DtypeObj, fill_value=np.nan):
  """
  Find the minimal dtype that can hold both the given dtype and fill_value.
 
  Parameters
  ----------
- dtype : np.dtype
+ dtype : np.dtype or ExtensionDtype
  fill_value : scalar, default np.nan
 
  Returns
@@ -593,9 +592,13 @@ def _maybe_promote_cached(dtype, fill_value, fill_value_type):
  return _maybe_promote(dtype, fill_value)
 
 
-def _maybe_promote(dtype: np.dtype, fill_value=np.nan):
+def _maybe_promote(dtype: DtypeObj, fill_value=np.nan):
  # The actual implementation of the function, use `maybe_promote` above for
  # a cached version.
+
+ if not isinstance(dtype, np.dtype):
+ return dtype._maybe_promote(fill_value)
+
  if not is_scalar(fill_value):
  # with object dtype there is nothing to promote, and the user can
  # pass pretty much any weird fill_value they like
@@ -611,12 +614,6 @@ def _maybe_promote(dtype: np.dtype, fill_value=np.nan):
  fv = na_value_for_dtype(dtype)
  return dtype, fv
 
- elif isinstance(dtype, CategoricalDtype):
- if fill_value in dtype.categories or isna(fill_value):
- return dtype, fill_value
- else:
- return object, ensure_object(fill_value)
-
  elif isna(fill_value):
  dtype = _dtype_obj
  if fill_value is None:

pandas/core/dtypes/dtypes.py

@@ -635,6 +635,15 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
 
  return find_common_type(non_cat_dtypes)
 
+ def _maybe_promote(self, item) -> tuple[DtypeObj, Any]:
+ from pandas.core.dtypes.missing import is_valid_na_for_dtype
+
+ if item in self.categories or is_valid_na_for_dtype(
+ item, self.categories.dtype
+ ):
+ return self, item
+ return np.dtype(object), item
+
 
 @register_extension_dtype
 class DatetimeTZDtype(PandasExtensionDtype):
@@ -1500,3 +1509,15 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
  return type(self).from_numpy_dtype(new_dtype)
  except (KeyError, NotImplementedError):
  return None
+
+ def _maybe_promote(self, item) -> tuple[DtypeObj, Any]:
+ from pandas.core.dtypes.cast import maybe_promote
+ from pandas.core.dtypes.missing import is_valid_na_for_dtype
+
+ if is_valid_na_for_dtype(item, self):
+ return self, item
+
+ dtype, item = maybe_promote(self.numpy_dtype, item)
+ if dtype.kind in "iufb":
+ return type(self).from_numpy_dtype(dtype), item
+ return dtype, item

pandas/core/indexing.py

@@ -2091,7 +2091,7 @@ def _setitem_with_indexer_missing(self, indexer, value):
  return self._setitem_with_indexer(new_indexer, value, "loc")
 
  # this preserves dtype of the value and of the object
- if not is_scalar(value):
+ if is_list_like(value):
  new_dtype = None
 
  elif is_valid_na_for_dtype(value, self.obj.dtype):
@@ -2107,8 +2107,7 @@ def _setitem_with_indexer_missing(self, indexer, value):
  # We should not cast, if we have object dtype because we can
  # set timedeltas into object series
  curr_dtype = self.obj.dtype
- curr_dtype = getattr(curr_dtype, "numpy_dtype", curr_dtype)
- new_dtype = maybe_promote(curr_dtype, value)[0]
+ new_dtype, value = maybe_promote(curr_dtype, value)
  else:
  new_dtype = None
 

pandas/tests/extension/test_arrow.py

@@ -2855,6 +2855,29 @@ def test_describe_timedelta_data(pa_type):
  tm.assert_series_equal(result, expected)
 
 
+@pytest.mark.parametrize(
+ "value, target_value, dtype",
+ [
+ (pa.scalar(4, type="int32"), 4, "int32[pyarrow]"),
+ (pa.scalar(4, type="int64"), 4, "int32[pyarrow]"),
+ # (pa.scalar(4.5, type="float64"), 4, "int32[pyarrow]"),
+ (4, 4, "int32[pyarrow]"),
+ (pd.NA, None, "int32[pyarrow]"),
+ (None, None, "int32[pyarrow]"),
+ (pa.scalar(None, type="int32"), None, "int32[pyarrow]"),
+ (pa.scalar(None, type="int64"), None, "int32[pyarrow]"),
+ ],
+)
+def test_series_setitem_with_enlargement(value, target_value, dtype):
+ # GH#52235
+ # similar to series/inedexing/test_setitem.py::test_setitem_keep_precision
+ # and test_setitem_enlarge_with_na, but for arrow dtypes
+ ser = pd.Series([1, 2, 3], dtype=dtype)
+ ser[3] = value
+ expected = pd.Series([1, 2, 3, target_value], dtype=dtype)
+ tm.assert_series_equal(ser, expected)
+
+
 @pytest.mark.parametrize("pa_type", tm.DATETIME_PYARROW_DTYPES)
 def test_describe_datetime_data(pa_type):
  # GH53001