Provide more precise type hints for itertools.permutations (python#11019

)
Avasam · Nov 11, 2023 · 6764465 · 6764465
1 parent 571cc6d
commit 6764465
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi
@@ -241,10 +241,19 @@ class product(Iterator[_T_co]):
     def __iter__(self) -> Self: ...
     def __next__(self) -> _T_co: ...
 
-class permutations(Iterator[tuple[_T, ...]], Generic[_T]):
-    def __init__(self, iterable: Iterable[_T], r: int | None = ...) -> None: ...
+class permutations(Iterator[_T_co]):
+    @overload
+    def __new__(cls, iterable: Iterable[_T], r: Literal[2]) -> permutations[tuple[_T, _T]]: ...
+    @overload
+    def __new__(cls, iterable: Iterable[_T], r: Literal[3]) -> permutations[tuple[_T, _T, _T]]: ...
+    @overload
+    def __new__(cls, iterable: Iterable[_T], r: Literal[4]) -> permutations[tuple[_T, _T, _T, _T]]: ...
+    @overload
+    def __new__(cls, iterable: Iterable[_T], r: Literal[5]) -> permutations[tuple[_T, _T, _T, _T, _T]]: ...
+    @overload
+    def __new__(cls, iterable: Iterable[_T], r: int | None = ...) -> permutations[tuple[_T, ...]]: ...
     def __iter__(self) -> Self: ...
-    def __next__(self) -> tuple[_T, ...]: ...
+    def __next__(self) -> _T_co: ...
 
 class combinations(Iterator[_T_co]):
     @overload