Skip to content

Commit 9d7f771

Browse files
committed
tqdm: add thread_map and process_map type hints
1 parent 33b3568 commit 9d7f771

1 file changed

Lines changed: 144 additions & 2 deletions

File tree

Lines changed: 144 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,146 @@
1+
from _typeshed import SupportsWrite
2+
from collections.abc import Callable, Iterable, Mapping
3+
from typing import Any, TypeVar, overload
4+
from typing_extensions import TypedDict, Unpack
5+
6+
from ..std import tqdm
7+
18
__all__ = ["thread_map", "process_map"]
29

3-
def thread_map(fn, *iterables, **tqdm_kwargs): ...
4-
def process_map(fn, *iterables, **tqdm_kwargs): ...
10+
_R = TypeVar("_R")
11+
_T1 = TypeVar("_T1")
12+
_T2 = TypeVar("_T2")
13+
_T3 = TypeVar("_T3")
14+
_T4 = TypeVar("_T4")
15+
_T5 = TypeVar("_T5")
16+
17+
class _TqdmKwargs(TypedDict, total=False):
18+
# Concurrent-specific parameters
19+
tqdm_class: type[tqdm[object]]
20+
max_workers: int | None
21+
chunksize: int
22+
lock_name: str
23+
# Standard tqdm parameters
24+
desc: str | None
25+
total: float | None
26+
leave: bool | None
27+
file: SupportsWrite[str] | None
28+
ncols: int | None
29+
mininterval: float
30+
maxinterval: float
31+
miniters: float | None
32+
ascii: bool | str | None
33+
disable: bool | None
34+
unit: str
35+
unit_scale: bool | float
36+
dynamic_ncols: bool
37+
smoothing: float
38+
bar_format: str | None
39+
initial: float
40+
position: int | None
41+
postfix: Mapping[str, object] | str | None
42+
unit_divisor: float
43+
write_bytes: bool | None
44+
lock_args: tuple[bool | None, float | None] | tuple[bool | None] | None
45+
nrows: int | None
46+
colour: str | None
47+
delay: float | None
48+
49+
@overload
50+
def thread_map(fn: Callable[[_T1], _R], iter1: Iterable[_T1], /, **tqdm_kwargs: Unpack[_TqdmKwargs]) -> list[_R]: ...
51+
@overload
52+
def thread_map(
53+
fn: Callable[[_T1, _T2], _R], iter1: Iterable[_T1], iter2: Iterable[_T2], /, **tqdm_kwargs: Unpack[_TqdmKwargs]
54+
) -> list[_R]: ...
55+
@overload
56+
def thread_map(
57+
fn: Callable[[_T1, _T2, _T3], _R],
58+
iter1: Iterable[_T1],
59+
iter2: Iterable[_T2],
60+
iter3: Iterable[_T3],
61+
/,
62+
**tqdm_kwargs: Unpack[_TqdmKwargs],
63+
) -> list[_R]: ...
64+
@overload
65+
def thread_map(
66+
fn: Callable[[_T1, _T2, _T3, _T4], _R],
67+
iter1: Iterable[_T1],
68+
iter2: Iterable[_T2],
69+
iter3: Iterable[_T3],
70+
iter4: Iterable[_T4],
71+
/,
72+
**tqdm_kwargs: Unpack[_TqdmKwargs],
73+
) -> list[_R]: ...
74+
@overload
75+
def thread_map(
76+
fn: Callable[[_T1, _T2, _T3, _T4, _T5], _R],
77+
iter1: Iterable[_T1],
78+
iter2: Iterable[_T2],
79+
iter3: Iterable[_T3],
80+
iter4: Iterable[_T4],
81+
iter5: Iterable[_T5],
82+
/,
83+
**tqdm_kwargs: Unpack[_TqdmKwargs],
84+
) -> list[_R]: ...
85+
@overload
86+
def thread_map(
87+
fn: Callable[..., _R],
88+
iter1: Iterable[Any],
89+
iter2: Iterable[Any],
90+
iter3: Iterable[Any],
91+
iter4: Iterable[Any],
92+
iter5: Iterable[Any],
93+
iter6: Iterable[Any],
94+
/,
95+
*iterables: Iterable[Any],
96+
**tqdm_kwargs: Unpack[_TqdmKwargs],
97+
) -> list[_R]: ...
98+
@overload
99+
def process_map(fn: Callable[[_T1], _R], iter1: Iterable[_T1], /, **tqdm_kwargs: Unpack[_TqdmKwargs]) -> list[_R]: ...
100+
@overload
101+
def process_map(
102+
fn: Callable[[_T1, _T2], _R], iter1: Iterable[_T1], iter2: Iterable[_T2], /, **tqdm_kwargs: Unpack[_TqdmKwargs]
103+
) -> list[_R]: ...
104+
@overload
105+
def process_map(
106+
fn: Callable[[_T1, _T2, _T3], _R],
107+
iter1: Iterable[_T1],
108+
iter2: Iterable[_T2],
109+
iter3: Iterable[_T3],
110+
/,
111+
**tqdm_kwargs: Unpack[_TqdmKwargs],
112+
) -> list[_R]: ...
113+
@overload
114+
def process_map(
115+
fn: Callable[[_T1, _T2, _T3, _T4], _R],
116+
iter1: Iterable[_T1],
117+
iter2: Iterable[_T2],
118+
iter3: Iterable[_T3],
119+
iter4: Iterable[_T4],
120+
/,
121+
**tqdm_kwargs: Unpack[_TqdmKwargs],
122+
) -> list[_R]: ...
123+
@overload
124+
def process_map(
125+
fn: Callable[[_T1, _T2, _T3, _T4, _T5], _R],
126+
iter1: Iterable[_T1],
127+
iter2: Iterable[_T2],
128+
iter3: Iterable[_T3],
129+
iter4: Iterable[_T4],
130+
iter5: Iterable[_T5],
131+
/,
132+
**tqdm_kwargs: Unpack[_TqdmKwargs],
133+
) -> list[_R]: ...
134+
@overload
135+
def process_map(
136+
fn: Callable[..., _R],
137+
iter1: Iterable[Any],
138+
iter2: Iterable[Any],
139+
iter3: Iterable[Any],
140+
iter4: Iterable[Any],
141+
iter5: Iterable[Any],
142+
iter6: Iterable[Any],
143+
/,
144+
*iterables: Iterable[Any],
145+
**tqdm_kwargs: Unpack[_TqdmKwargs],
146+
) -> list[_R]: ...

0 commit comments

Comments
 (0)