[tqdm] Add thread_map and process_map type hints (#15460)

Author: hoel-bagard

stubs/tqdm/tqdm/contrib/concurrent.pyi

@@ -1,4 +1,139 @@
+from _typeshed import SupportsWrite
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any, TypedDict, TypeVar, overload, type_check_only
+from typing_extensions import Unpack
+
+from ..std import tqdm
+
 __all__ = ["thread_map", "process_map"]
 
-def thread_map(fn, *iterables, **tqdm_kwargs): ...
-def process_map(fn, *iterables, **tqdm_kwargs): ...
+_R = TypeVar("_R")
+_T1 = TypeVar("_T1")
+_T2 = TypeVar("_T2")
+_T3 = TypeVar("_T3")
+_T4 = TypeVar("_T4")
+_T5 = TypeVar("_T5")
+
+@type_check_only
+class _TqdmKwargs(TypedDict, total=False):
+    # Concurrent-specific parameters
+    tqdm_class: type[tqdm[object]]
+    max_workers: int | None
+    chunksize: int
+    lock_name: str
+    # Standard tqdm parameters
+    desc: str | None
+    total: float | None
+    leave: bool | None
+    file: SupportsWrite[str] | None
+    ncols: int | None
+    mininterval: float
+    maxinterval: float
+    miniters: float | None
+    ascii: bool | str | None
+    disable: bool | None
+    unit: str
+    unit_scale: bool | float
+    dynamic_ncols: bool
+    smoothing: float
+    bar_format: str | None
+    initial: float
+    position: int | None
+    postfix: Mapping[str, object] | str | None
+    unit_divisor: float
+    write_bytes: bool | None
+    lock_args: tuple[bool | None, float | None] | tuple[bool | None] | None
+    nrows: int | None
+    colour: str | None
+    delay: float | None
+
+@overload
+def thread_map(fn: Callable[[_T1], _R], iter1: Iterable[_T1], **tqdm_kwargs: Unpack[_TqdmKwargs]) -> list[_R]: ...
+@overload
+def thread_map(
+    fn: Callable[[_T1, _T2], _R], iter1: Iterable[_T1], iter2: Iterable[_T2], /, **tqdm_kwargs: Unpack[_TqdmKwargs]
+) -> list[_R]: ...
+@overload
+def thread_map(
+    fn: Callable[[_T1, _T2, _T3], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def thread_map(
+    fn: Callable[[_T1, _T2, _T3, _T4], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    iter4: Iterable[_T4],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def thread_map(
+    fn: Callable[[_T1, _T2, _T3, _T4, _T5], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    iter4: Iterable[_T4],
+    iter5: Iterable[_T5],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def thread_map(
+    fn: Callable[..., _R],
+    iter1: Iterable[Any],
+    iter2: Iterable[Any],
+    iter3: Iterable[Any],
+    iter4: Iterable[Any],
+    iter5: Iterable[Any],
+    iter6: Iterable[Any],
+    *iterables: Iterable[Any],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def process_map(fn: Callable[[_T1], _R], iter1: Iterable[_T1], **tqdm_kwargs: Unpack[_TqdmKwargs]) -> list[_R]: ...
+@overload
+def process_map(
+    fn: Callable[[_T1, _T2], _R], iter1: Iterable[_T1], iter2: Iterable[_T2], **tqdm_kwargs: Unpack[_TqdmKwargs]
+) -> list[_R]: ...
+@overload
+def process_map(
+    fn: Callable[[_T1, _T2, _T3], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def process_map(
+    fn: Callable[[_T1, _T2, _T3, _T4], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    iter4: Iterable[_T4],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def process_map(
+    fn: Callable[[_T1, _T2, _T3, _T4, _T5], _R],
+    iter1: Iterable[_T1],
+    iter2: Iterable[_T2],
+    iter3: Iterable[_T3],
+    iter4: Iterable[_T4],
+    iter5: Iterable[_T5],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...
+@overload
+def process_map(
+    fn: Callable[..., _R],
+    iter1: Iterable[Any],
+    iter2: Iterable[Any],
+    iter3: Iterable[Any],
+    iter4: Iterable[Any],
+    iter5: Iterable[Any],
+    iter6: Iterable[Any],
+    *iterables: Iterable[Any],
+    **tqdm_kwargs: Unpack[_TqdmKwargs],
+) -> list[_R]: ...