add extension test

Author: jbrockmendel

pandas/tests/extension/base/setitem.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from pandas.core.dtypes.common import is_hashable
+
 import pandas as pd
 import pandas._testing as tm
 
@@ -310,6 +312,22 @@ def test_setitem_expand_with_extension(self, data):
         result.loc[:, "B"] = data
         tm.assert_frame_equal(result, expected)
 
+    def test_loc_setitem_with_expansion_preserves_ea_index_dtype(self, data):
+        # GH#41626 retain index.dtype in setitem-with-expansion
+        if not is_hashable(data[0]):
+            pytest.skip("Test does not apply to non-hashable data.")
+        data = data.unique()
+        expected = pd.DataFrame({"A": range(len(data))}, index=data)
+        df = expected.iloc[:-1]
+        ser = df["A"]
+        item = data[-1]
+
+        df.loc[item] = len(data) - 1
+        tm.assert_frame_equal(df, expected)
+
+        ser.loc[item] = len(data) - 1
+        tm.assert_series_equal(ser, expected["A"])
+
     def test_setitem_frame_invalid_length(self, data):
         df = pd.DataFrame({"A": [1] * len(data)})
         xpr = (

pandas/tests/extension/test_arrow.py

@@ -1067,6 +1067,15 @@ def test_comp_masked_numpy(self, masked_dtype, comparison_op):
         expected = pd.Series(exp, dtype=ArrowDtype(pa.bool_()))
         tm.assert_series_equal(result, expected)
 
+    def test_loc_setitem_with_expansion_preserves_ea_index_dtype(self, data, request):
+        pa_dtype = data.dtype.pyarrow_dtype
+        if pa.types.is_date(pa_dtype):
+            mark = pytest.mark.xfail(
+                reason="GH#62343 incorrectly casts to timestamp[ms][pyarrow]"
+            )
+            request.applymarker(mark)
+        super().test_loc_setitem_with_expansion_preserves_ea_index_dtype(data)
+
 
 class TestLogicalOps:
     """Various Series and DataFrame logical ops methods."""

pandas/tests/extension/test_interval.py

@@ -126,6 +126,13 @@ def test_EA_types(self, engine, data, request):
     def test_astype_str(self, data):
         super().test_astype_str(data)
 
+    @pytest.mark.xfail(
+        reason="Test is invalid for IntervalDtype, needs to be adapted for "
+        "this dtype with an index with index._index_as_unique."
+    )
+    def test_loc_setitem_with_expansion_preserves_ea_index_dtype(self, data):
+        super().test_loc_setitem_with_expansion_preserves_ea_index_dtype(data)
+
 
 # TODO: either belongs in tests.arrays.interval or move into base tests.
 def test_fillna_non_scalar_raises(data_missing):

pandas/tests/extension/test_masked.py

@@ -360,3 +360,9 @@ def check_accumulate(self, ser: pd.Series, op_name: str, skipna: bool):
             )
         )
         tm.assert_series_equal(result, expected)
+
+    def test_loc_setitem_with_expansion_preserves_ea_index_dtype(self, data, request):
+        if data.dtype.kind == "b":
+            mark = pytest.mark.xfail(reason="GH#62344 incorrectly casts to object")
+            request.applymarker(mark)
+        super().test_loc_setitem_with_expansion_preserves_ea_index_dtype(data)

pandas/tests/extension/test_numpy.py

@@ -421,6 +421,12 @@ def test_index_from_listlike_with_dtype(self, data):
     def test_EA_types(self, engine, data, request):
         super().test_EA_types(engine, data, request)
 
+    def test_loc_setitem_with_expansion_preserves_ea_index_dtype(self, data, request):
+        if isinstance(data[-1], tuple):
+            mark = pytest.mark.xfail(reason="Unpacks tuple")
+            request.applymarker(mark)
+        super().test_loc_setitem_with_expansion_preserves_ea_index_dtype(data)
+
 
 class Test2DCompat(base.NDArrayBacked2DTests):
     pass