Skip to content

Commit 8611eca

Browse files
authored
TYP: misc return types (#57430)
1 parent 0863b22 commit 8611eca

File tree

6 files changed

+281
-62
lines changed

6 files changed

+281
-62
lines changed

pandas/_typing.py

+1
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@
118118
Concatenate: Any = None
119119

120120
HashableT = TypeVar("HashableT", bound=Hashable)
121+
HashableT2 = TypeVar("HashableT2", bound=Hashable)
121122
MutableMappingT = TypeVar("MutableMappingT", bound=MutableMapping)
122123

123124
# array-like

pandas/core/base.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
from typing import (
99
TYPE_CHECKING,
1010
Any,
11-
Callable,
1211
Generic,
1312
Literal,
1413
cast,
@@ -105,7 +104,7 @@ class PandasObject(DirNamesMixin):
105104
_cache: dict[str, Any]
106105

107106
@property
108-
def _constructor(self) -> Callable[..., Self]:
107+
def _constructor(self) -> type[Self]:
109108
"""
110109
Class constructor (for this class it's just `__class__`).
111110
"""
@@ -1356,7 +1355,7 @@ def searchsorted(
13561355
sorter=sorter,
13571356
)
13581357

1359-
def drop_duplicates(self, *, keep: DropKeep = "first"):
1358+
def drop_duplicates(self, *, keep: DropKeep = "first") -> Self:
13601359
duplicated = self._duplicated(keep=keep)
13611360
# error: Value of type "IndexOpsMixin" is not indexable
13621361
return self[~duplicated] # type: ignore[index]

pandas/core/frame.py

+105-17
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,8 @@
217217
FormattersType,
218218
Frequency,
219219
FromDictOrient,
220+
HashableT,
221+
HashableT2,
220222
IgnoreRaise,
221223
IndexKeyFunc,
222224
IndexLabel,
@@ -239,6 +241,7 @@
239241
SortKind,
240242
StorageOptions,
241243
Suffixes,
244+
T,
242245
ToStataByteorder,
243246
ToTimestampHow,
244247
UpdateJoin,
@@ -643,10 +646,10 @@ class DataFrame(NDFrame, OpsMixin):
643646
__pandas_priority__ = 4000
644647

645648
@property
646-
def _constructor(self) -> Callable[..., DataFrame]:
649+
def _constructor(self) -> type[DataFrame]:
647650
return DataFrame
648651

649-
def _constructor_from_mgr(self, mgr, axes):
652+
def _constructor_from_mgr(self, mgr, axes) -> DataFrame:
650653
if self._constructor is DataFrame:
651654
# we are pandas.DataFrame (or a subclass that doesn't override _constructor)
652655
return DataFrame._from_mgr(mgr, axes=axes)
@@ -659,7 +662,7 @@ def _constructor_from_mgr(self, mgr, axes):
659662
def _sliced_from_mgr(self, mgr, axes) -> Series:
660663
return Series._from_mgr(mgr, axes)
661664

662-
def _constructor_sliced_from_mgr(self, mgr, axes):
665+
def _constructor_sliced_from_mgr(self, mgr, axes) -> Series:
663666
if self._constructor_sliced is Series:
664667
ser = self._sliced_from_mgr(mgr, axes)
665668
ser._name = None # caller is responsible for setting real name
@@ -1353,7 +1356,7 @@ def _get_values_for_csv(
13531356
decimal: str,
13541357
na_rep: str,
13551358
quoting, # int csv.QUOTE_FOO from stdlib
1356-
) -> Self:
1359+
) -> DataFrame:
13571360
# helper used by to_csv
13581361
mgr = self._mgr.get_values_for_csv(
13591362
float_format=float_format,
@@ -1831,7 +1834,7 @@ def from_dict(
18311834
a b 1 3
18321835
c 2 4
18331836
"""
1834-
index = None
1837+
index: list | Index | None = None
18351838
orient = orient.lower() # type: ignore[assignment]
18361839
if orient == "index":
18371840
if len(data) > 0:
@@ -1857,7 +1860,7 @@ def from_dict(
18571860
else:
18581861
realdata = data["data"]
18591862

1860-
def create_index(indexlist, namelist):
1863+
def create_index(indexlist, namelist) -> Index:
18611864
index: Index
18621865
if len(namelist) > 1:
18631866
index = MultiIndex.from_tuples(indexlist, names=namelist)
@@ -2700,6 +2703,42 @@ def to_feather(self, path: FilePath | WriteBuffer[bytes], **kwargs) -> None:
27002703

27012704
to_feather(self, path, **kwargs)
27022705

2706+
@overload
2707+
def to_markdown(
2708+
self,
2709+
buf: None = ...,
2710+
*,
2711+
mode: str = ...,
2712+
index: bool = ...,
2713+
storage_options: StorageOptions | None = ...,
2714+
**kwargs,
2715+
) -> str:
2716+
...
2717+
2718+
@overload
2719+
def to_markdown(
2720+
self,
2721+
buf: FilePath | WriteBuffer[str],
2722+
*,
2723+
mode: str = ...,
2724+
index: bool = ...,
2725+
storage_options: StorageOptions | None = ...,
2726+
**kwargs,
2727+
) -> None:
2728+
...
2729+
2730+
@overload
2731+
def to_markdown(
2732+
self,
2733+
buf: FilePath | WriteBuffer[str] | None,
2734+
*,
2735+
mode: str = ...,
2736+
index: bool = ...,
2737+
storage_options: StorageOptions | None = ...,
2738+
**kwargs,
2739+
) -> str | None:
2740+
...
2741+
27032742
@doc(
27042743
Series.to_markdown,
27052744
klass=_shared_doc_kwargs["klass"],
@@ -2881,6 +2920,39 @@ def to_parquet(
28812920
**kwargs,
28822921
)
28832922

2923+
@overload
2924+
def to_orc(
2925+
self,
2926+
path: None = ...,
2927+
*,
2928+
engine: Literal["pyarrow"] = ...,
2929+
index: bool | None = ...,
2930+
engine_kwargs: dict[str, Any] | None = ...,
2931+
) -> bytes:
2932+
...
2933+
2934+
@overload
2935+
def to_orc(
2936+
self,
2937+
path: FilePath | WriteBuffer[bytes],
2938+
*,
2939+
engine: Literal["pyarrow"] = ...,
2940+
index: bool | None = ...,
2941+
engine_kwargs: dict[str, Any] | None = ...,
2942+
) -> None:
2943+
...
2944+
2945+
@overload
2946+
def to_orc(
2947+
self,
2948+
path: FilePath | WriteBuffer[bytes] | None,
2949+
*,
2950+
engine: Literal["pyarrow"] = ...,
2951+
index: bool | None = ...,
2952+
engine_kwargs: dict[str, Any] | None = ...,
2953+
) -> bytes | None:
2954+
...
2955+
28842956
def to_orc(
28852957
self,
28862958
path: FilePath | WriteBuffer[bytes] | None = None,
@@ -4027,7 +4099,7 @@ def _setitem_slice(self, key: slice, value) -> None:
40274099
# backwards-compat, xref GH#31469
40284100
self.iloc[key] = value
40294101

4030-
def _setitem_array(self, key, value):
4102+
def _setitem_array(self, key, value) -> None:
40314103
# also raises Exception if object array with NA values
40324104
if com.is_bool_indexer(key):
40334105
# bool indexer is indexing along rows
@@ -4061,7 +4133,7 @@ def _setitem_array(self, key, value):
40614133
elif np.ndim(value) > 1:
40624134
# list of lists
40634135
value = DataFrame(value).values
4064-
return self._setitem_array(key, value)
4136+
self._setitem_array(key, value)
40654137

40664138
else:
40674139
self._iset_not_inplace(key, value)
@@ -4595,7 +4667,7 @@ def eval(self, expr: str, *, inplace: bool = False, **kwargs) -> Any | None:
45954667

45964668
return _eval(expr, inplace=inplace, **kwargs)
45974669

4598-
def select_dtypes(self, include=None, exclude=None) -> Self:
4670+
def select_dtypes(self, include=None, exclude=None) -> DataFrame:
45994671
"""
46004672
Return a subset of the DataFrame's columns based on the column dtypes.
46014673
@@ -5474,9 +5546,21 @@ def pop(self, item: Hashable) -> Series:
54745546
"""
54755547
return super().pop(item=item)
54765548

5549+
@overload
5550+
def _replace_columnwise(
5551+
self, mapping: dict[Hashable, tuple[Any, Any]], inplace: Literal[True], regex
5552+
) -> None:
5553+
...
5554+
5555+
@overload
5556+
def _replace_columnwise(
5557+
self, mapping: dict[Hashable, tuple[Any, Any]], inplace: Literal[False], regex
5558+
) -> Self:
5559+
...
5560+
54775561
def _replace_columnwise(
54785562
self, mapping: dict[Hashable, tuple[Any, Any]], inplace: bool, regex
5479-
):
5563+
) -> Self | None:
54805564
"""
54815565
Dispatch to Series.replace column-wise.
54825566
@@ -5505,7 +5589,7 @@ def _replace_columnwise(
55055589
res._iset_item(i, newobj, inplace=inplace)
55065590

55075591
if inplace:
5508-
return
5592+
return None
55095593
return res.__finalize__(self)
55105594

55115595
@doc(NDFrame.shift, klass=_shared_doc_kwargs["klass"])
@@ -11815,19 +11899,19 @@ def kurt(
1181511899
product = prod
1181611900

1181711901
@doc(make_doc("cummin", ndim=2))
11818-
def cummin(self, axis: Axis | None = None, skipna: bool = True, *args, **kwargs):
11902+
def cummin(self, axis: Axis = 0, skipna: bool = True, *args, **kwargs) -> Self:
1181911903
return NDFrame.cummin(self, axis, skipna, *args, **kwargs)
1182011904

1182111905
@doc(make_doc("cummax", ndim=2))
11822-
def cummax(self, axis: Axis | None = None, skipna: bool = True, *args, **kwargs):
11906+
def cummax(self, axis: Axis = 0, skipna: bool = True, *args, **kwargs) -> Self:
1182311907
return NDFrame.cummax(self, axis, skipna, *args, **kwargs)
1182411908

1182511909
@doc(make_doc("cumsum", ndim=2))
11826-
def cumsum(self, axis: Axis | None = None, skipna: bool = True, *args, **kwargs):
11910+
def cumsum(self, axis: Axis = 0, skipna: bool = True, *args, **kwargs) -> Self:
1182711911
return NDFrame.cumsum(self, axis, skipna, *args, **kwargs)
1182811912

1182911913
@doc(make_doc("cumprod", 2))
11830-
def cumprod(self, axis: Axis | None = None, skipna: bool = True, *args, **kwargs):
11914+
def cumprod(self, axis: Axis = 0, skipna: bool = True, *args, **kwargs) -> Self:
1183111915
return NDFrame.cumprod(self, axis, skipna, *args, **kwargs)
1183211916

1183311917
def nunique(self, axis: Axis = 0, dropna: bool = True) -> Series:
@@ -12710,8 +12794,12 @@ def values(self) -> np.ndarray:
1271012794
return self._mgr.as_array()
1271112795

1271212796

12713-
def _from_nested_dict(data) -> collections.defaultdict:
12714-
new_data: collections.defaultdict = collections.defaultdict(dict)
12797+
def _from_nested_dict(
12798+
data: Mapping[HashableT, Mapping[HashableT2, T]],
12799+
) -> collections.defaultdict[HashableT2, dict[HashableT, T]]:
12800+
new_data: collections.defaultdict[
12801+
HashableT2, dict[HashableT, T]
12802+
] = collections.defaultdict(dict)
1271512803
for index, s in data.items():
1271612804
for col, v in s.items():
1271712805
new_data[col][index] = v

0 commit comments

Comments
 (0)