[heapq] Use overloads to prevent unsound uses (#15669)

Author: jonathandung

stdlib/_heapq.pyi

@@ -4,7 +4,7 @@ from typing import Final
 
 __about__: Final[str]
 
-def heapify(heap: list[_T], /) -> None: ...
+def heapify(heap: list[_T], /) -> None: ...  # To work around the fact that list is invariant
 def heappop(heap: list[_T], /) -> _T: ...
 def heappush(heap: list[_T], item: _T, /) -> None: ...
 def heappushpop(heap: list[_T], item: _T, /) -> _T: ...

stdlib/heapq.pyi

@@ -1,8 +1,8 @@
 import sys
 from _heapq import *
-from _typeshed import SupportsRichComparison
+from _typeshed import SupportsRichComparison, SupportsRichComparisonT as _T
 from collections.abc import Callable, Generator, Iterable
-from typing import Any, Final, TypeVar
+from typing import Final, TypeVar, overload
 
 __all__ = ["heappush", "heappop", "heapify", "heapreplace", "merge", "nlargest", "nsmallest", "heappushpop"]
 
@@ -14,9 +14,19 @@ _S = TypeVar("_S")
 
 __about__: Final[str]
 
-def merge(
-    *iterables: Iterable[_S], key: Callable[[_S], SupportsRichComparison] | None = None, reverse: bool = False
-) -> Generator[_S]: ...
-def nlargest(n: int, iterable: Iterable[_S], key: Callable[[_S], SupportsRichComparison] | None = None) -> list[_S]: ...
-def nsmallest(n: int, iterable: Iterable[_S], key: Callable[[_S], SupportsRichComparison] | None = None) -> list[_S]: ...
-def _heapify_max(heap: list[Any], /) -> None: ...  # undocumented
+@overload
+def merge(*iterables: Iterable[_S], key: Callable[[_S], SupportsRichComparison], reverse: bool = False) -> Generator[_S]: ...
+@overload
+def merge(*iterables: Iterable[_T], key: None = None, reverse: bool = False) -> Generator[_T]: ...
+
+@overload
+def nlargest(n: int, iterable: Iterable[_S], key: Callable[[_S], SupportsRichComparison]) -> list[_S]: ...
+@overload
+def nlargest(n: int, iterable: Iterable[_T], key: None = None) -> list[_T]: ...
+
+@overload
+def nsmallest(n: int, iterable: Iterable[_S], key: Callable[[_S], SupportsRichComparison]) -> list[_S]: ...
+@overload
+def nsmallest(n: int, iterable: Iterable[_T], key: None = None) -> list[_T]: ...
+
+def _heapify_max(heap: list[SupportsRichComparison], /) -> None: ...  # undocumented