From ed6dc176da089c17b74cbc62811e2e82373efdb9 Mon Sep 17 00:00:00 2001
From: Matthias Kramm <kramm@quiss.org>
Date: Mon, 12 Jun 2017 20:10:08 -0700
Subject: [PATCH] make itertools.product more precise (#1393)

---
 stdlib/2/itertools.pyi | 38 +++++++++++++++++++++++++++++---
 stdlib/3/itertools.pyi | 50 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/stdlib/2/itertools.pyi b/stdlib/2/itertools.pyi
index 4d9edbc32..284047763 100644
--- a/stdlib/2/itertools.pyi
+++ b/stdlib/2/itertools.pyi
@@ -93,9 +93,41 @@ def izip(iter1: Iterable[Any], iter2: Iterable[Any],
 def izip_longest(*p: Iterable[Any],
                  fillvalue: Any = ...) -> Iterator[Any]: ...
 
-# TODO: Return type should be Iterator[Tuple[..]], but unknown tuple shape.
-#       Iterator[Sequence[_T]] loses this type information.
-def product(*p: Iterable[_T], repeat: int = ...) -> Iterator[Sequence[_T]]: ...
+@overload
+def product(iter1: Iterable[_T1]) -> Iterator[Tuple[_T1]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2]) -> Iterator[Tuple[_T1, _T2]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3]) -> Iterator[Tuple[_T1, _T2, _T3]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4]) -> Iterator[Tuple[_T1, _T2, _T3, _T4]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4],
+            iter5: Iterable[_T5]) -> Iterator[Tuple[_T1, _T2, _T3, _T4, _T5]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4],
+            iter5: Iterable[_T5],
+            iter6: Iterable[_T6]) -> Iterator[Tuple[_T1, _T2, _T3, _T4, _T5, _T6]]: ...
+@overload
+def product(iter1: Iterable[Any],
+            iter2: Iterable[Any],
+            iter3: Iterable[Any],
+            iter4: Iterable[Any],
+            iter5: Iterable[Any],
+            iter6: Iterable[Any],
+            iter7: Iterable[Any], *iterables: Iterable) -> Iterator[Tuple]: ...
 
 def permutations(iterable: Iterable[_T],
                  r: int = ...) -> Iterator[Sequence[_T]]: ...
diff --git a/stdlib/3/itertools.pyi b/stdlib/3/itertools.pyi
index 3e4379a03..1f79cd257 100644
--- a/stdlib/3/itertools.pyi
+++ b/stdlib/3/itertools.pyi
@@ -51,7 +51,55 @@ def tee(iterable: Iterable[_T], n: int = ...) -> Tuple[Iterator[_T], ...]: ...
 def zip_longest(*p: Iterable[Any],
                 fillvalue: Any = ...) -> Iterator[Any]: ...
 
-def product(*p: Iterable[_T], repeat: int = ...) -> Iterator[Tuple[_T, ...]]: ...
+_T1 = TypeVar('_T1')
+_T2 = TypeVar('_T2')
+_T3 = TypeVar('_T3')
+_T4 = TypeVar('_T4')
+_T5 = TypeVar('_T5')
+_T6 = TypeVar('_T6')
+
+@overload
+def product(iter1: Iterable[_T1], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1, _T2]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1, _T2, _T3]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1, _T2, _T3, _T4]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4],
+            iter5: Iterable[_T5], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1, _T2, _T3, _T4, _T5]]: ...
+@overload
+def product(iter1: Iterable[_T1],
+            iter2: Iterable[_T2],
+            iter3: Iterable[_T3],
+            iter4: Iterable[_T4],
+            iter5: Iterable[_T5],
+            iter6: Iterable[_T6], *,
+            repeat: int = ...) -> Iterator[Tuple[_T1, _T2, _T3, _T4, _T5, _T6]]: ...
+@overload
+def product(iter1: Iterable[Any],
+            iter2: Iterable[Any],
+            iter3: Iterable[Any],
+            iter4: Iterable[Any],
+            iter5: Iterable[Any],
+            iter6: Iterable[Any],
+            iter7: Iterable[Any], *iterables: Iterable,
+            repeat: int = ...) -> Iterator[Tuple]: ...
 
 def permutations(iterable: Iterable[_T],
                  r: Optional[int] = ...) -> Iterator[Tuple[_T, ...]]: ...