BUG: remove usage of _constructor._get_axis_number (fix when _constructor properties return callables) (#46018)

Author: jorisvandenbossche

doc/source/whatsnew/v1.4.2.rst

@@ -23,6 +23,7 @@ Fixed regressions
 
 Bug fixes
 ~~~~~~~~~
+- Fix some cases for subclasses that define their ``_constructor`` properties as general callables (:issue:`46018`)
 - Fixed "longtable" formatting in :meth:`.Styler.to_latex` when ``column_format`` is given in extended format (:issue:`46037`)
 -
 

pandas/_testing/__init__.py

@@ -803,23 +803,28 @@ class SubclassedSeries(Series):
 
     @property
     def _constructor(self):
-        return SubclassedSeries
+        # For testing, those properties return a generic callable, and not
+        # the actual class. In this case that is equivalent, but it is to
+        # ensure we don't rely on the property returning a class
+        # See https://github.com/pandas-dev/pandas/pull/46018 and
+        # https://github.com/pandas-dev/pandas/issues/32638 and linked issues
+        return lambda *args, **kwargs: SubclassedSeries(*args, **kwargs)
 
     @property
     def _constructor_expanddim(self):
-        return SubclassedDataFrame
+        return lambda *args, **kwargs: SubclassedDataFrame(*args, **kwargs)
 
 
 class SubclassedDataFrame(DataFrame):
     _metadata = ["testattr"]
 
     @property
     def _constructor(self):
-        return SubclassedDataFrame
+        return lambda *args, **kwargs: SubclassedDataFrame(*args, **kwargs)
 
     @property
     def _constructor_sliced(self):
-        return SubclassedSeries
+        return lambda *args, **kwargs: SubclassedSeries(*args, **kwargs)
 
 
 class SubclassedCategorical(Categorical):

pandas/core/frame.py

@@ -581,10 +581,10 @@ class DataFrame(NDFrame, OpsMixin):
     _mgr: BlockManager | ArrayManager
 
     @property
-    def _constructor(self) -> type[DataFrame]:
+    def _constructor(self) -> Callable[..., DataFrame]:
         return DataFrame
 
-    _constructor_sliced: type[Series] = Series
+    _constructor_sliced: Callable[..., Series] = Series
 
     # ----------------------------------------------------------------------
     # Constructors

pandas/core/generic.py

@@ -444,7 +444,7 @@ def _validate_dtype(cls, dtype) -> DtypeObj | None:
     # Construction
 
     @property
-    def _constructor(self: NDFrameT) -> type[NDFrameT]:
+    def _constructor(self: NDFrameT) -> Callable[..., NDFrameT]:
         """
         Used when a manipulation result has the same dimensions as the
         original.
@@ -779,17 +779,9 @@ def swapaxes(self: NDFrameT, axis1, axis2, copy=True) -> NDFrameT:
         if copy:
             new_values = new_values.copy()
 
-        # ignore needed because of NDFrame constructor is different than
-        # DataFrame/Series constructors.
         return self._constructor(
-            # error: Argument 1 to "NDFrame" has incompatible type "ndarray"; expected
-            # "Union[ArrayManager, BlockManager]"
-            # error: Argument 2 to "NDFrame" has incompatible type "*Generator[Index,
-            # None, None]"; expected "bool" [arg-type]
-            # error: Argument 2 to "NDFrame" has incompatible type "*Generator[Index,
-            # None, None]"; expected "Optional[Mapping[Hashable, Any]]"
-            new_values,  # type: ignore[arg-type]
-            *new_axes,  # type: ignore[arg-type]
+            new_values,
+            *new_axes,
         ).__finalize__(self, method="swapaxes")
 
     @final
@@ -2088,11 +2080,7 @@ def __array_wrap__(
             # ptp also requires the item_from_zerodim
             return res
         d = self._construct_axes_dict(self._AXIS_ORDERS, copy=False)
-        # error: Argument 1 to "NDFrame" has incompatible type "ndarray";
-        # expected "BlockManager"
-        return self._constructor(res, **d).__finalize__(  # type: ignore[arg-type]
-            self, method="__array_wrap__"
-        )
+        return self._constructor(res, **d).__finalize__(self, method="__array_wrap__")
 
     @final
     def __array_ufunc__(
@@ -5923,11 +5911,9 @@ def astype(
         # GH 19920: retain column metadata after concat
         result = concat(results, axis=1, copy=False)
         # GH#40810 retain subclass
-        # Incompatible types in assignment (expression has type "NDFrameT",
-        # variable has type "DataFrame")
-        # Argument 1 to "NDFrame" has incompatible type "DataFrame"; expected
-        # "Union[ArrayManager, SingleArrayManager, BlockManager, SingleBlockManager]"
-        result = self._constructor(result)  # type: ignore[arg-type,assignment]
+        # error: Incompatible types in assignment
+        # (expression has type "NDFrameT", variable has type "DataFrame")
+        result = self._constructor(result)  # type: ignore[assignment]
         result.columns = self.columns
         result = result.__finalize__(self, method="astype")
         # https://github.com/python/mypy/issues/8354
@@ -6613,8 +6599,10 @@ def replace(
 
             if isinstance(to_replace, (tuple, list)):
                 if isinstance(self, ABCDataFrame):
+                    from pandas import Series
+
                     result = self.apply(
-                        self._constructor_sliced._replace_single,
+                        Series._replace_single,
                         args=(to_replace, method, inplace, limit),
                     )
                     if inplace:
@@ -9139,11 +9127,7 @@ def _where(
 
             # we are the same shape, so create an actual object for alignment
             else:
-                # error: Argument 1 to "NDFrame" has incompatible type "ndarray";
-                # expected "BlockManager"
-                other = self._constructor(
-                    other, **self._construct_axes_dict()  # type: ignore[arg-type]
-                )
+                other = self._constructor(other, **self._construct_axes_dict())
 
         if axis is None:
             axis = 0

pandas/core/groupby/groupby.py

@@ -1733,7 +1733,7 @@ def _cumcount_array(self, ascending: bool = True) -> np.ndarray:
 
     @final
     @property
-    def _obj_1d_constructor(self) -> type[Series]:
+    def _obj_1d_constructor(self) -> Callable:
         # GH28330 preserve subclassed Series/DataFrames
         if isinstance(self.obj, DataFrame):
             return self.obj._constructor_sliced
@@ -2158,7 +2158,7 @@ def size(self) -> DataFrame | Series:
             )
 
         # GH28330 preserve subclassed Series/DataFrames through calls
-        if issubclass(self.obj._constructor, Series):
+        if isinstance(self.obj, Series):
             result = self._obj_1d_constructor(result, name=self.obj.name)
         else:
             result = self._obj_1d_constructor(result)

pandas/core/reshape/concat.py

@@ -6,6 +6,7 @@
 from collections import abc
 from typing import (
     TYPE_CHECKING,
+    Callable,
     Hashable,
     Iterable,
     Literal,
@@ -467,7 +468,9 @@ def __init__(
 
         # Standardize axis parameter to int
         if isinstance(sample, ABCSeries):
-            axis = sample._constructor_expanddim._get_axis_number(axis)
+            from pandas import DataFrame
+
+            axis = DataFrame._get_axis_number(axis)
         else:
             axis = sample._get_axis_number(axis)
 
@@ -539,7 +542,7 @@ def __init__(
         self.new_axes = self._get_new_axes()
 
     def get_result(self):
-        cons: type[DataFrame | Series]
+        cons: Callable[..., DataFrame | Series]
         sample: DataFrame | Series
 
         # series only

pandas/core/series.py

@@ -526,11 +526,11 @@ def _init_dict(
     # ----------------------------------------------------------------------
 
     @property
-    def _constructor(self) -> type[Series]:
+    def _constructor(self) -> Callable[..., Series]:
         return Series
 
     @property
-    def _constructor_expanddim(self) -> type[DataFrame]:
+    def _constructor_expanddim(self) -> Callable[..., DataFrame]:
         """
         Used when a manipulation result has one higher dimension as the
         original, such as Series.to_frame()

pandas/tests/frame/test_subclass.py

@@ -737,3 +737,11 @@ def test_equals_subclass(self):
         df2 = tm.SubclassedDataFrame({"a": [1, 2, 3]})
         assert df1.equals(df2)
         assert df2.equals(df1)
+
+    def test_replace_list_method(self):
+        # https://github.com/pandas-dev/pandas/pull/46018
+        df = tm.SubclassedDataFrame({"A": [0, 1, 2]})
+        result = df.replace([1, 2], method="ffill")
+        expected = tm.SubclassedDataFrame({"A": [0, 0, 0]})
+        assert isinstance(result, tm.SubclassedDataFrame)
+        tm.assert_frame_equal(result, expected)

pandas/tests/groupby/test_groupby_subclass.py

@@ -46,7 +46,7 @@ def test_groupby_preserves_subclass(obj, groupby_func):
     # Reduction or transformation kernels should preserve type
     slices = {"ngroup", "cumcount", "size"}
     if isinstance(obj, DataFrame) and groupby_func in slices:
-        assert isinstance(result1, obj._constructor_sliced)
+        assert isinstance(result1, tm.SubclassedSeries)
     else:
         assert isinstance(result1, type(obj))