From 40215d1fa3ef33a31fe68624472c10dd62350f18 Mon Sep 17 00:00:00 2001
From: Yannack <yannack@users.noreply.github.com>
Date: Thu, 18 Jul 2019 15:38:28 -0400
Subject: [PATCH] Fix the definition of nsmallest() in stdlib/2 and
 2and3/heapq.pyi (#3114)

Missing support of the optional "key" kwarg in nsmallest.
Also fixed nlargest syntax for 2and3 which was also missing.
Use a protocol for heapq.py
---
 stdlib/2/heapq.pyi      | 10 +++++++---
 stdlib/2and3/_heapq.pyi |  8 +++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/stdlib/2/heapq.pyi b/stdlib/2/heapq.pyi
index 00abb31a6..488221c7c 100644
--- a/stdlib/2/heapq.pyi
+++ b/stdlib/2/heapq.pyi
@@ -1,7 +1,10 @@
-from typing import TypeVar, List, Iterable, Any, Callable, Optional
+from typing import TypeVar, List, Iterable, Any, Callable, Optional, Protocol
 
 _T = TypeVar('_T')
 
+class _Sortable(Protocol):
+    def __lt__(self: _T, other: _T) -> bool: ...
+
 def cmp_lt(x, y) -> bool: ...
 def heappush(heap: List[_T], item: _T) -> None: ...
 def heappop(heap: List[_T]) -> _T:
@@ -12,5 +15,6 @@ def heapreplace(heap: List[_T], item: _T) -> _T:
     raise IndexError()  # if heap is empty
 def merge(*iterables: Iterable[_T]) -> Iterable[_T]: ...
 def nlargest(n: int, iterable: Iterable[_T],
-             key: Optional[Callable[[_T], Any]] = ...) -> List[_T]: ...
-def nsmallest(n: int, iterable: Iterable[_T]) -> List[_T]: ...
+             key: Optional[Callable[[_T], _Sortable]] = ...) -> List[_T]: ...
+def nsmallest(n: int, iterable: Iterable[_T],
+              key: Optional[Callable[[_T], _Sortable]] = ...) -> List[_T]: ...
diff --git a/stdlib/2and3/_heapq.pyi b/stdlib/2and3/_heapq.pyi
index 8b7f6eac0..9ff4a08fe 100644
--- a/stdlib/2and3/_heapq.pyi
+++ b/stdlib/2and3/_heapq.pyi
@@ -1,6 +1,7 @@
 """Stub file for the '_heapq' module."""
 
-from typing import TypeVar, List
+from typing import TypeVar, List, Iterable, Any, Callable, Optional
+import sys
 
 _T = TypeVar("_T")
 
@@ -11,5 +12,6 @@ def heappush(heap: List[_T], item: _T) -> None: ...
 def heappushpop(heap: List[_T], item: _T) -> _T: ...
 def heapreplace(heap: List[_T], item: _T) -> _T:
     raise IndexError()  # if list is empty
-def nlargest(a: int, b: List[_T]) -> List[_T]: ...
-def nsmallest(a: int, b: List[_T]) -> List[_T]: ...
+if sys.version_info < (3,):
+    def nlargest(n: int, iterable: Iterable[_T]) -> List[_T]: ...
+    def nsmallest(n: int, iterable: Iterable[_T]) -> List[_T]: ...