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

Co-authored-by: Matthew Roeschke <emailformattr@gmail.com>

Author: simonjayhawkins

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):