Skip to content

Commit 5344107

Browse files
authored
TYP: Axis-related annotations (#48907)
* TYP: Axis-related annotations * black
1 parent 6e95d7a commit 5344107

File tree

7 files changed

+27
-39
lines changed

7 files changed

+27
-39
lines changed

pandas/core/apply.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,7 @@ def agg_dict_like(self) -> DataFrame | Series:
506506
ktu._set_names(selected_obj.columns.names)
507507
keys_to_use = ktu
508508

509-
axis = 0 if isinstance(obj, ABCSeries) else 1
509+
axis: AxisInt = 0 if isinstance(obj, ABCSeries) else 1
510510
# error: Key expression in dictionary comprehension has incompatible type
511511
# "Hashable"; expected type "NDFrame" [misc]
512512
result = concat(
@@ -932,7 +932,7 @@ def apply_str(self) -> DataFrame | Series:
932932

933933

934934
class FrameRowApply(FrameApply):
935-
axis = 0
935+
axis: AxisInt = 0
936936

937937
def apply_broadcast(self, target: DataFrame) -> DataFrame:
938938
return super().apply_broadcast(target)
@@ -992,7 +992,7 @@ def wrap_results_for_axis(
992992

993993

994994
class FrameColumnApply(FrameApply):
995-
axis = 1
995+
axis: AxisInt = 1
996996

997997
def apply_broadcast(self, target: DataFrame) -> DataFrame:
998998
result = super().apply_broadcast(target.T)
@@ -1069,7 +1069,7 @@ def infer_to_same_shape(self, results: ResType, res_index: Index) -> DataFrame:
10691069

10701070
class SeriesApply(NDFrameApply):
10711071
obj: Series
1072-
axis = 0
1072+
axis: AxisInt = 0
10731073

10741074
def __init__(
10751075
self,
@@ -1199,7 +1199,7 @@ def transform(self):
11991199

12001200

12011201
class ResamplerWindowApply(Apply):
1202-
axis = 0
1202+
axis: AxisInt = 0
12031203
obj: Resampler | BaseWindow
12041204

12051205
def __init__(

pandas/core/frame.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -7586,7 +7586,7 @@ class diet
75867586
# Arithmetic Methods
75877587

75887588
def _cmp_method(self, other, op):
7589-
axis = 1 # only relevant for Series other case
7589+
axis: Literal[1] = 1 # only relevant for Series other case
75907590

75917591
self, other = ops.align_method_FRAME(self, other, axis, flex=False, level=None)
75927592

@@ -7598,7 +7598,7 @@ def _arith_method(self, other, op):
75987598
if ops.should_reindex_frame_op(self, other, op, 1, 1, None, None):
75997599
return ops.frame_arith_method_with_reindex(self, other, op)
76007600

7601-
axis = 1 # only relevant for Series other case
7601+
axis: Literal[1] = 1 # only relevant for Series other case
76027602
other = ops.maybe_prepare_scalar_for_op(other, (self.shape[axis],))
76037603

76047604
self, other = ops.align_method_FRAME(self, other, axis, flex=True, level=None)
@@ -10657,7 +10657,7 @@ def c(x):
1065710657
if not drop:
1065810658
# Find non-matching labels along the given axis
1065910659
# and append missing correlations (GH 22375)
10660-
raxis = 1 if axis == 0 else 0
10660+
raxis: AxisInt = 1 if axis == 0 else 0
1066110661
result_index = this._get_axis(raxis).union(other._get_axis(raxis))
1066210662
idx_diff = result_index.difference(correl.index)
1066310663

@@ -11666,7 +11666,7 @@ def isin(self, values: Series | DataFrame | Sequence | Mapping) -> DataFrame:
1166611666
"columns": 1,
1166711667
}
1166811668
_AXIS_LEN = len(_AXIS_ORDERS)
11669-
_info_axis_number = 1
11669+
_info_axis_number: Literal[1] = 1
1167011670
_info_axis_name: Literal["columns"] = "columns"
1167111671

1167211672
index = properties.AxisProperty(

pandas/core/generic.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -577,7 +577,7 @@ def _get_axis(self, axis: Axis) -> Index:
577577

578578
@final
579579
@classmethod
580-
def _get_block_manager_axis(cls, axis: Axis) -> int:
580+
def _get_block_manager_axis(cls, axis: Axis) -> AxisInt:
581581
"""Map the axis to the block_manager axis."""
582582
axis = cls._get_axis_number(axis)
583583
ndim = cls._AXIS_LEN
@@ -1705,7 +1705,7 @@ def _is_level_reference(self, key: Level, axis: Axis = 0) -> bool_t:
17051705
)
17061706

17071707
@final
1708-
def _is_label_reference(self, key: Level, axis: int = 0) -> bool_t:
1708+
def _is_label_reference(self, key: Level, axis: Axis = 0) -> bool_t:
17091709
"""
17101710
Test whether a key is a label reference for a given axis.
17111711
@@ -1725,8 +1725,8 @@ def _is_label_reference(self, key: Level, axis: int = 0) -> bool_t:
17251725
-------
17261726
is_label: bool
17271727
"""
1728-
axis = self._get_axis_number(axis)
1729-
other_axes = (ax for ax in range(self._AXIS_LEN) if ax != axis)
1728+
axis_int = self._get_axis_number(axis)
1729+
other_axes = (ax for ax in range(self._AXIS_LEN) if ax != axis_int)
17301730

17311731
return (
17321732
key is not None

pandas/core/indexing.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1111,7 +1111,7 @@ class _LocIndexer(_LocationIndexer):
11111111
# Key Checks
11121112

11131113
@doc(_LocationIndexer._validate_key)
1114-
def _validate_key(self, key, axis: AxisInt):
1114+
def _validate_key(self, key, axis: Axis):
11151115
# valid for a collection of labels (we check their presence later)
11161116
# slice of labels (where start-end in labels)
11171117
# slice of integers (only if in the labels)

pandas/core/reshape/concat.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from pandas._typing import (
2323
Axis,
24+
AxisInt,
2425
HashableT,
2526
)
2627
from pandas.util._decorators import (
@@ -636,7 +637,7 @@ def _get_new_axes(self) -> list[Index]:
636637
for i in range(ndim)
637638
]
638639

639-
def _get_comb_axis(self, i: int) -> Index:
640+
def _get_comb_axis(self, i: AxisInt) -> Index:
640641
data_axis = self.objs[0]._get_block_manager_axis(i)
641642
return get_objs_combined_axis(
642643
self.objs,

pandas/core/series.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -6228,7 +6228,7 @@ def mask( # type: ignore[override]
62286228
# Add index
62296229
_AXIS_ORDERS: list[Literal["index", "columns"]] = ["index"]
62306230
_AXIS_LEN = len(_AXIS_ORDERS)
6231-
_info_axis_number = 0
6231+
_info_axis_number: Literal[0] = 0
62326232
_info_axis_name: Literal["index"] = "index"
62336233

62346234
index = properties.AxisProperty(

pandas/io/pytables.py

+10-23
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def _ensure_decoded(s):
137137
return s
138138

139139

140-
def _ensure_encoding(encoding):
140+
def _ensure_encoding(encoding: str | None) -> str:
141141
# set the encoding if we need
142142
if encoding is None:
143143
encoding = _default_encoding
@@ -543,8 +543,6 @@ class HDFStore:
543543

544544
_handle: File | None
545545
_mode: str
546-
_complevel: int
547-
_fletcher32: bool
548546

549547
def __init__(
550548
self,
@@ -1043,7 +1041,7 @@ def select_as_multiple(
10431041
_tbls = [x for x in tbls if isinstance(x, Table)]
10441042

10451043
# axis is the concentration axes
1046-
axis = list({t.non_index_axes[0][0] for t in _tbls})[0]
1044+
axis = {t.non_index_axes[0][0] for t in _tbls}.pop()
10471045

10481046
def func(_start, _stop, _where):
10491047

@@ -1975,9 +1973,6 @@ class IndexCol:
19751973
is_data_indexable: bool = True
19761974
_info_fields = ["freq", "tz", "index_name"]
19771975

1978-
name: str
1979-
cname: str
1980-
19811976
def __init__(
19821977
self,
19831978
name: str,
@@ -2310,7 +2305,7 @@ def __init__(
23102305
values=None,
23112306
kind=None,
23122307
typ=None,
2313-
cname=None,
2308+
cname: str | None = None,
23142309
pos=None,
23152310
tz=None,
23162311
ordered=None,
@@ -2622,17 +2617,14 @@ class Fixed:
26222617
format_type: str = "fixed" # GH#30962 needed by dask
26232618
obj_type: type[DataFrame | Series]
26242619
ndim: int
2625-
encoding: str
26262620
parent: HDFStore
2627-
group: Node
2628-
errors: str
26292621
is_table: bool = False
26302622

26312623
def __init__(
26322624
self,
26332625
parent: HDFStore,
26342626
group: Node,
2635-
encoding: str = "UTF-8",
2627+
encoding: str | None = "UTF-8",
26362628
errors: str = "strict",
26372629
) -> None:
26382630
assert isinstance(parent, HDFStore), type(parent)
@@ -3306,24 +3298,19 @@ class Table(Fixed):
33063298
levels: int | list[Hashable] = 1
33073299
is_table = True
33083300

3309-
index_axes: list[IndexCol]
3310-
non_index_axes: list[tuple[int, Any]]
3311-
values_axes: list[DataCol]
3312-
data_columns: list
33133301
metadata: list
3314-
info: dict
33153302

33163303
def __init__(
33173304
self,
33183305
parent: HDFStore,
33193306
group: Node,
3320-
encoding=None,
3307+
encoding: str | None = None,
33213308
errors: str = "strict",
3322-
index_axes=None,
3323-
non_index_axes=None,
3324-
values_axes=None,
3325-
data_columns=None,
3326-
info=None,
3309+
index_axes: list[IndexCol] | None = None,
3310+
non_index_axes: list[tuple[AxisInt, Any]] | None = None,
3311+
values_axes: list[DataCol] | None = None,
3312+
data_columns: list | None = None,
3313+
info: dict | None = None,
33273314
nan_rep=None,
33283315
) -> None:
33293316
super().__init__(parent, group, encoding=encoding, errors=errors)

0 commit comments

Comments
 (0)