TYP: setter for index/columns property-like (AxisProperty) (pandas-dev#46565)

Co-authored-by: Matthew Roeschke <[email protected]>

Author: 2 people

pandas/_libs/properties.pyi

@@ -1,9 +1,28 @@
-# pyright: reportIncompleteStub = false
-from typing import Any
+from typing import (
+    Sequence,
+    overload,
+)
+
+from pandas._typing import (
+    AnyArrayLike,
+    DataFrame,
+    Index,
+    Series,
+)
 
 # note: this is a lie to make type checkers happy (they special
 # case property). cache_readonly uses attribute names similar to
 # property (fget) but it does not provide fset and fdel.
 cache_readonly = property
 
-def __getattr__(name: str) -> Any: ...  # incomplete
+class AxisProperty:
+
+    axis: int
+    def __init__(self, axis: int = ..., doc: str = ...) -> None: ...
+    @overload
+    def __get__(self, obj: DataFrame | Series, type) -> Index: ...
+    @overload
+    def __get__(self, obj: None, type) -> AxisProperty: ...
+    def __set__(
+        self, obj: DataFrame | Series, value: AnyArrayLike | Sequence
+    ) -> None: ...

pandas/core/apply.py

@@ -235,8 +235,12 @@ def transform(self) -> DataFrame | Series:
             and not obj.empty
         ):
             raise ValueError("Transform function failed")
+        # error: Argument 1 to "__get__" of "AxisProperty" has incompatible type
+        # "Union[Series, DataFrame, GroupBy[Any], SeriesGroupBy,
+        # DataFrameGroupBy, BaseWindow, Resampler]"; expected "Union[DataFrame,
+        # Series]"
         if not isinstance(result, (ABCSeries, ABCDataFrame)) or not result.index.equals(
-            obj.index
+            obj.index  # type:ignore[arg-type]
         ):
             raise ValueError("Function did not transform")
 
@@ -645,7 +649,11 @@ class NDFrameApply(Apply):
 
     @property
     def index(self) -> Index:
-        return self.obj.index
+        # error: Argument 1 to "__get__" of "AxisProperty" has incompatible type
+        # "Union[Series, DataFrame, GroupBy[Any], SeriesGroupBy,
+        # DataFrameGroupBy, BaseWindow, Resampler]"; expected "Union[DataFrame,
+        # Series]"
+        return self.obj.index  # type:ignore[arg-type]
 
     @property
     def agg_axis(self) -> Index:

pandas/core/frame.py

@@ -11199,12 +11199,10 @@ def isin(self, values) -> DataFrame:
     _info_axis_number = 1
     _info_axis_name = "columns"
 
-    index: Index = properties.AxisProperty(
+    index = properties.AxisProperty(
         axis=1, doc="The index (row labels) of the DataFrame."
     )
-    columns: Index = properties.AxisProperty(
-        axis=0, doc="The column labels of the DataFrame."
-    )
+    columns = properties.AxisProperty(axis=0, doc="The column labels of the DataFrame.")
 
     @property
     def _AXIS_NUMBERS(self) -> dict[str, int]:

pandas/core/generic.py

@@ -37,6 +37,7 @@
     to_offset,
 )
 from pandas._typing import (
+    AnyArrayLike,
     ArrayLike,
     Axis,
     CompressionOptions,
@@ -757,7 +758,7 @@ def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t):
             obj.set_axis(labels, axis=axis, inplace=True)
             return obj
 
-    def _set_axis(self, axis: int, labels: Index) -> None:
+    def _set_axis(self, axis: int, labels: AnyArrayLike | Sequence) -> None:
         labels = ensure_index(labels)
         self._mgr.set_axis(axis, labels)
         self._clear_item_cache()

pandas/core/groupby/generic.py

@@ -275,9 +275,9 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
             func = maybe_mangle_lambdas(func)
             ret = self._aggregate_multiple_funcs(func)
             if relabeling:
-                # error: Incompatible types in assignment (expression has type
-                # "Optional[List[str]]", variable has type "Index")
-                ret.columns = columns  # type: ignore[assignment]
+                # columns is not narrowed by mypy from relabeling flag
+                assert columns is not None  # for mypy
+                ret.columns = columns
             return ret
 
         else:

pandas/core/series.py

@@ -33,6 +33,7 @@
 from pandas._libs.lib import no_default
 from pandas._typing import (
     AggFuncType,
+    AnyArrayLike,
     ArrayLike,
     Axis,
     Dtype,
@@ -553,7 +554,7 @@ def _constructor_expanddim(self) -> Callable[..., DataFrame]:
     def _can_hold_na(self) -> bool:
         return self._mgr._can_hold_na
 
-    def _set_axis(self, axis: int, labels) -> None:
+    def _set_axis(self, axis: int, labels: AnyArrayLike | Sequence) -> None:
         """
         Override generic, we want to set the _typ here.
 
@@ -5813,7 +5814,7 @@ def mask(
     _info_axis_number = 0
     _info_axis_name = "index"
 
-    index: Index = properties.AxisProperty(
+    index = properties.AxisProperty(
         axis=0, doc="The index (axis labels) of the Series."
     )
 

pandas/io/parsers/arrow_parser_wrapper.py

@@ -105,12 +105,7 @@ def _finalize_output(self, frame: DataFrame) -> DataFrame:
                 multi_index_named = False
             frame.columns = self.names
         # we only need the frame not the names
-        # error: Incompatible types in assignment (expression has type
-        # "Union[List[Union[Union[str, int, float, bool], Union[Period, Timestamp,
-        # Timedelta, Any]]], Index]", variable has type "Index")  [assignment]
-        frame.columns, frame = self._do_date_conversions(  # type: ignore[assignment]
-            frame.columns, frame
-        )
+        frame.columns, frame = self._do_date_conversions(frame.columns, frame)
         if self.index_col is not None:
             for i, item in enumerate(self.index_col):
                 if is_integer(item):