From c0e9c4f4c55b293e88e68e5b7977750ee851728d Mon Sep 17 00:00:00 2001 From: Steve B Date: Fri, 3 Jul 2020 15:23:43 +0200 Subject: [PATCH] add overloads for combinations r 2-5 (#4309) This doesn't cover everything but it means it a lot of common usages the type will be slightly more useful Co-authored-by: steve brazier --- stdlib/3/itertools.pyi | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/stdlib/3/itertools.pyi b/stdlib/3/itertools.pyi index e9965e885..d930050f1 100644 --- a/stdlib/3/itertools.pyi +++ b/stdlib/3/itertools.pyi @@ -4,6 +4,7 @@ import sys from typing import Any, Callable, Generic, Iterable, Iterator, Optional, Tuple, TypeVar, overload +from typing_extensions import Literal _T = TypeVar("_T") _S = TypeVar("_S") @@ -97,5 +98,14 @@ def product( @overload def product(*iterables: Iterable[Any], repeat: int = ...) -> Iterator[Tuple[Any, ...]]: ... def permutations(iterable: Iterable[_T], r: Optional[int] = ...) -> Iterator[Tuple[_T, ...]]: ... +@overload +def combinations(iterable: Iterable[_T], r: Literal[2]) -> Iterator[Tuple[_T, _T]]: ... +@overload +def combinations(iterable: Iterable[_T], r: Literal[3]) -> Iterator[Tuple[_T, _T, _T]]: ... +@overload +def combinations(iterable: Iterable[_T], r: Literal[4]) -> Iterator[Tuple[_T, _T, _T, _T]]: ... +@overload +def combinations(iterable: Iterable[_T], r: Literal[5]) -> Iterator[Tuple[_T, _T, _T, _T, _T]]: ... +@overload def combinations(iterable: Iterable[_T], r: int) -> Iterator[Tuple[_T, ...]]: ... def combinations_with_replacement(iterable: Iterable[_T], r: int) -> Iterator[Tuple[_T, ...]]: ...