Skip to content

Commit a1f9ae2

Browse files
authored
TYP: annotations for internals, set_axis (#32376)
1 parent 8b347ea commit a1f9ae2

File tree

4 files changed

+27
-23
lines changed

4 files changed

+27
-23
lines changed

pandas/core/frame.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -3599,7 +3599,7 @@ def align(
35993599
see_also_sub=" or columns",
36003600
)
36013601
@Appender(NDFrame.set_axis.__doc__)
3602-
def set_axis(self, labels, axis=0, inplace=False):
3602+
def set_axis(self, labels, axis: Axis = 0, inplace: bool = False):
36033603
return super().set_axis(labels, axis=axis, inplace=inplace)
36043604

36053605
@Substitution(**_shared_doc_kwargs)

pandas/core/generic.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -520,7 +520,7 @@ def _obj_with_exclusions(self: FrameOrSeries) -> FrameOrSeries:
520520
""" internal compat with SelectionMixin """
521521
return self
522522

523-
def set_axis(self, labels, axis=0, inplace=False):
523+
def set_axis(self, labels, axis: Axis = 0, inplace: bool = False):
524524
"""
525525
Assign desired index to given axis.
526526
@@ -561,7 +561,8 @@ def set_axis(self, labels, axis=0, inplace=False):
561561
obj.set_axis(labels, axis=axis, inplace=True)
562562
return obj
563563

564-
def _set_axis(self, axis, labels) -> None:
564+
def _set_axis(self, axis: int, labels: Index) -> None:
565+
labels = ensure_index(labels)
565566
self._data.set_axis(axis, labels)
566567
self._clear_item_cache()
567568

pandas/core/internals/managers.py

+10-10
Original file line numberDiff line numberDiff line change
@@ -164,15 +164,15 @@ def __nonzero__(self):
164164
__bool__ = __nonzero__
165165

166166
@property
167-
def shape(self):
167+
def shape(self) -> Tuple[int, ...]:
168168
return tuple(len(ax) for ax in self.axes)
169169

170170
@property
171171
def ndim(self) -> int:
172172
return len(self.axes)
173173

174-
def set_axis(self, axis, new_labels):
175-
new_labels = ensure_index(new_labels)
174+
def set_axis(self, axis: int, new_labels: Index):
175+
# Caller is responsible for ensuring we have an Index object.
176176
old_len = len(self.axes[axis])
177177
new_len = len(new_labels)
178178

@@ -184,7 +184,9 @@ def set_axis(self, axis, new_labels):
184184

185185
self.axes[axis] = new_labels
186186

187-
def rename_axis(self, mapper, axis, copy: bool = True, level=None):
187+
def rename_axis(
188+
self, mapper, axis: int, copy: bool = True, level=None
189+
) -> "BlockManager":
188190
"""
189191
Rename one of axes.
190192
@@ -193,7 +195,7 @@ def rename_axis(self, mapper, axis, copy: bool = True, level=None):
193195
mapper : unary callable
194196
axis : int
195197
copy : bool, default True
196-
level : int, default None
198+
level : int or None, default None
197199
"""
198200
obj = self.copy(deep=copy)
199201
obj.set_axis(axis, _transform_index(self.axes[axis], mapper, level))
@@ -233,7 +235,7 @@ def _rebuild_blknos_and_blklocs(self):
233235
self._blklocs = new_blklocs
234236

235237
@property
236-
def items(self):
238+
def items(self) -> Index:
237239
return self.axes[0]
238240

239241
def _get_counts(self, f):
@@ -623,7 +625,7 @@ def comp(s, regex=False):
623625
bm._consolidate_inplace()
624626
return bm
625627

626-
def is_consolidated(self):
628+
def is_consolidated(self) -> bool:
627629
"""
628630
Return True if more than one block with the same dtype
629631
"""
@@ -688,7 +690,7 @@ def get_numeric_data(self, copy: bool = False):
688690
self._consolidate_inplace()
689691
return self.combine([b for b in self.blocks if b.is_numeric], copy)
690692

691-
def combine(self, blocks, copy=True):
693+
def combine(self, blocks: List[Block], copy: bool = True) -> "BlockManager":
692694
""" return a new manager with the blocks """
693695
if len(blocks) == 0:
694696
return self.make_empty()
@@ -992,7 +994,6 @@ def delete(self, item):
992994
self.blocks = tuple(
993995
b for blkno, b in enumerate(self.blocks) if not is_blk_deleted[blkno]
994996
)
995-
self._shape = None
996997
self._rebuild_blknos_and_blklocs()
997998

998999
def set(self, item, value):
@@ -1160,7 +1161,6 @@ def insert(self, loc: int, item, value, allow_duplicates: bool = False):
11601161

11611162
self.axes[0] = new_axis
11621163
self.blocks += (block,)
1163-
self._shape = None
11641164

11651165
self._known_consolidated = False
11661166

pandas/core/series.py

+13-10
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
TYPE_CHECKING,
1010
Any,
1111
Callable,
12-
Hashable,
1312
Iterable,
1413
List,
1514
Optional,
@@ -23,7 +22,7 @@
2322
from pandas._config import get_option
2423

2524
from pandas._libs import lib, properties, reshape, tslibs
26-
from pandas._typing import Label
25+
from pandas._typing import Axis, DtypeObj, Label
2726
from pandas.compat.numpy import function as nv
2827
from pandas.util._decorators import Appender, Substitution, doc
2928
from pandas.util._validators import validate_bool_kwarg, validate_percentile
@@ -177,7 +176,7 @@ class Series(base.IndexOpsMixin, generic.NDFrame):
177176

178177
_typ = "series"
179178

180-
_name: Optional[Hashable]
179+
_name: Label
181180
_metadata: List[str] = ["name"]
182181
_internal_names_set = {"index"} | generic.NDFrame._internal_names_set
183182
_accessors = {"dt", "cat", "str", "sparse"}
@@ -391,9 +390,12 @@ def _can_hold_na(self):
391390

392391
_index = None
393392

394-
def _set_axis(self, axis, labels, fastpath: bool = False) -> None:
393+
def _set_axis(self, axis: int, labels, fastpath: bool = False) -> None:
395394
"""
396395
Override generic, we want to set the _typ here.
396+
397+
This is called from the cython code when we set the `index` attribute
398+
directly, e.g. `series.index = [1, 2, 3]`.
397399
"""
398400
if not fastpath:
399401
labels = ensure_index(labels)
@@ -413,6 +415,7 @@ def _set_axis(self, axis, labels, fastpath: bool = False) -> None:
413415

414416
object.__setattr__(self, "_index", labels)
415417
if not fastpath:
418+
# The ensure_index call aabove ensures we have an Index object
416419
self._data.set_axis(axis, labels)
417420

418421
def _update_inplace(self, result, **kwargs):
@@ -421,25 +424,25 @@ def _update_inplace(self, result, **kwargs):
421424

422425
# ndarray compatibility
423426
@property
424-
def dtype(self):
427+
def dtype(self) -> DtypeObj:
425428
"""
426429
Return the dtype object of the underlying data.
427430
"""
428431
return self._data.dtype
429432

430433
@property
431-
def dtypes(self):
434+
def dtypes(self) -> DtypeObj:
432435
"""
433436
Return the dtype object of the underlying data.
434437
"""
435438
return self._data.dtype
436439

437440
@property
438-
def name(self) -> Optional[Hashable]:
441+
def name(self) -> Label:
439442
return self._name
440443

441444
@name.setter
442-
def name(self, value: Optional[Hashable]) -> None:
445+
def name(self, value: Label) -> None:
443446
if not is_hashable(value):
444447
raise TypeError("Series.name must be a hashable type")
445448
object.__setattr__(self, "_name", value)
@@ -689,7 +692,7 @@ def __array_ufunc__(
689692
inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs)
690693
result = getattr(ufunc, method)(*inputs, **kwargs)
691694

692-
name: Optional[Hashable]
695+
name: Label
693696
if len(set(names)) == 1:
694697
name = names[0]
695698
else:
@@ -3983,7 +3986,7 @@ def rename(
39833986
see_also_sub="",
39843987
)
39853988
@Appender(generic.NDFrame.set_axis.__doc__)
3986-
def set_axis(self, labels, axis=0, inplace=False):
3989+
def set_axis(self, labels, axis: Axis = 0, inplace: bool = False):
39873990
return super().set_axis(labels, axis=axis, inplace=inplace)
39883991

39893992
@Substitution(**_shared_doc_kwargs)

0 commit comments

Comments
 (0)