diff --git a/stubs/tqdm/tqdm/contrib/concurrent.pyi b/stubs/tqdm/tqdm/contrib/concurrent.pyi index 40e070a53c1c..c11da84b93ab 100644 --- a/stubs/tqdm/tqdm/contrib/concurrent.pyi +++ b/stubs/tqdm/tqdm/contrib/concurrent.pyi @@ -1,4 +1,138 @@ +from _typeshed import SupportsWrite +from collections.abc import Callable, Iterable, Mapping +from typing import Any, TypedDict, TypeVar, overload +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") + +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]: ...