Skip to content

Commit d16df29

Browse files
authored
TYP: Add cast to ABC Index-like types (#38043)
1 parent 2d661a8 commit d16df29

File tree

5 files changed

+60
-22
lines changed

5 files changed

+60
-22
lines changed

pandas/core/algorithms.py

+1
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ def _ensure_data(
168168
elif is_categorical_dtype(values.dtype) and (
169169
is_categorical_dtype(dtype) or dtype is None
170170
):
171+
values = cast("Categorical", values)
171172
values = values.codes
172173
dtype = pandas_dtype("category")
173174

pandas/core/dtypes/generic.py

+53-18
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,20 @@
44
from typing import TYPE_CHECKING, Type, cast
55

66
if TYPE_CHECKING:
7-
from pandas import DataFrame, Series
7+
from pandas import (
8+
CategoricalIndex,
9+
DataFrame,
10+
DatetimeIndex,
11+
Float64Index,
12+
Int64Index,
13+
IntervalIndex,
14+
MultiIndex,
15+
PeriodIndex,
16+
RangeIndex,
17+
Series,
18+
TimedeltaIndex,
19+
UInt64Index,
20+
)
821
from pandas.core.generic import NDFrame
922

1023

@@ -23,23 +36,45 @@ def _check(cls, inst) -> bool:
2336
return meta(name, (), dct)
2437

2538

26-
ABCInt64Index = create_pandas_abc_type("ABCInt64Index", "_typ", ("int64index",))
27-
ABCUInt64Index = create_pandas_abc_type("ABCUInt64Index", "_typ", ("uint64index",))
28-
ABCRangeIndex = create_pandas_abc_type("ABCRangeIndex", "_typ", ("rangeindex",))
29-
ABCFloat64Index = create_pandas_abc_type("ABCFloat64Index", "_typ", ("float64index",))
30-
ABCMultiIndex = create_pandas_abc_type("ABCMultiIndex", "_typ", ("multiindex",))
31-
ABCDatetimeIndex = create_pandas_abc_type(
32-
"ABCDatetimeIndex", "_typ", ("datetimeindex",)
33-
)
34-
ABCTimedeltaIndex = create_pandas_abc_type(
35-
"ABCTimedeltaIndex", "_typ", ("timedeltaindex",)
36-
)
37-
ABCPeriodIndex = create_pandas_abc_type("ABCPeriodIndex", "_typ", ("periodindex",))
38-
ABCCategoricalIndex = create_pandas_abc_type(
39-
"ABCCategoricalIndex", "_typ", ("categoricalindex",)
40-
)
41-
ABCIntervalIndex = create_pandas_abc_type(
42-
"ABCIntervalIndex", "_typ", ("intervalindex",)
39+
ABCInt64Index = cast(
40+
"Type[Int64Index]",
41+
create_pandas_abc_type("ABCInt64Index", "_typ", ("int64index",)),
42+
)
43+
ABCUInt64Index = cast(
44+
"Type[UInt64Index]",
45+
create_pandas_abc_type("ABCUInt64Index", "_typ", ("uint64index",)),
46+
)
47+
ABCRangeIndex = cast(
48+
"Type[RangeIndex]",
49+
create_pandas_abc_type("ABCRangeIndex", "_typ", ("rangeindex",)),
50+
)
51+
ABCFloat64Index = cast(
52+
"Type[Float64Index]",
53+
create_pandas_abc_type("ABCFloat64Index", "_typ", ("float64index",)),
54+
)
55+
ABCMultiIndex = cast(
56+
"Type[MultiIndex]",
57+
create_pandas_abc_type("ABCMultiIndex", "_typ", ("multiindex",)),
58+
)
59+
ABCDatetimeIndex = cast(
60+
"Type[DatetimeIndex]",
61+
create_pandas_abc_type("ABCDatetimeIndex", "_typ", ("datetimeindex",)),
62+
)
63+
ABCTimedeltaIndex = cast(
64+
"Type[TimedeltaIndex]",
65+
create_pandas_abc_type("ABCTimedeltaIndex", "_typ", ("timedeltaindex",)),
66+
)
67+
ABCPeriodIndex = cast(
68+
"Type[PeriodIndex]",
69+
create_pandas_abc_type("ABCPeriodIndex", "_typ", ("periodindex",)),
70+
)
71+
ABCCategoricalIndex = cast(
72+
"Type[CategoricalIndex]",
73+
create_pandas_abc_type("ABCCategoricalIndex", "_typ", ("categoricalindex",)),
74+
)
75+
ABCIntervalIndex = cast(
76+
"Type[IntervalIndex]",
77+
create_pandas_abc_type("ABCIntervalIndex", "_typ", ("intervalindex",)),
4378
)
4479
ABCIndexClass = create_pandas_abc_type(
4580
"ABCIndexClass",

pandas/core/indexes/datetimes.py

+1
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ class DatetimeIndex(DatetimeTimedeltaMixin):
227227
_is_numeric_dtype = False
228228

229229
_data: DatetimeArray
230+
inferred_freq: Optional[str]
230231
tz: Optional[tzinfo]
231232

232233
# --------------------------------------------------------------------

pandas/io/formats/latex.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -153,11 +153,11 @@ def pad_empties(x):
153153
break
154154
return [x[0]] + [i if i else " " * len(pad) for i in x[1:]]
155155

156-
out = (pad_empties(i) for i in out)
156+
gen = (pad_empties(i) for i in out)
157157

158158
# Add empty spaces for each column level
159159
clevels = self.frame.columns.nlevels
160-
out = [[" " * len(i[-1])] * clevels + i for i in out]
160+
out = [[" " * len(i[-1])] * clevels + i for i in gen]
161161

162162
# Add the column names to the last index column
163163
cnames = self.frame.columns.names

pandas/plotting/_matplotlib/timeseries.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# TODO: Use the fact that axis can have units to simplify the process
22

33
import functools
4-
from typing import TYPE_CHECKING, Optional
4+
from typing import TYPE_CHECKING, Optional, cast
55

66
import numpy as np
77

@@ -26,7 +26,7 @@
2626
if TYPE_CHECKING:
2727
from matplotlib.axes import Axes
2828

29-
from pandas import Index, Series
29+
from pandas import DatetimeIndex, Index, Series
3030

3131
# ---------------------------------------------------------------------
3232
# Plotting functions and monkey patches
@@ -243,6 +243,7 @@ def maybe_convert_index(ax: "Axes", data):
243243

244244
if freq is None:
245245
# We only get here for DatetimeIndex
246+
data.index = cast("DatetimeIndex", data.index)
246247
freq = data.index.inferred_freq
247248
freq = to_offset(freq)
248249

0 commit comments

Comments
 (0)