Improve accuracy of six byte index methods (#9117)

Author: Avasam

stdlib/_operator.pyi

@@ -1,5 +1,6 @@
 import sys
-from collections.abc import Callable, Container, Iterable, Mapping, MutableMapping, MutableSequence, Sequence
+from _typeshed import SupportsGetItem
+from collections.abc import Callable, Container, Iterable, MutableMapping, MutableSequence, Sequence
 from typing import Any, AnyStr, Generic, Protocol, SupportsAbs, TypeVar, overload
 from typing_extensions import ParamSpec, SupportsIndex, TypeAlias, final
 
@@ -77,11 +78,9 @@ def delitem(__a: MutableSequence[Any], __b: slice) -> None: ...
 @overload
 def delitem(__a: MutableMapping[_K, Any], __b: _K) -> None: ...
 @overload
-def getitem(__a: Sequence[_T], __b: SupportsIndex) -> _T: ...
-@overload
 def getitem(__a: Sequence[_T], __b: slice) -> Sequence[_T]: ...
 @overload
-def getitem(__a: Mapping[_K, _V], __b: _K) -> _V: ...
+def getitem(__a: SupportsGetItem[_K, _V], __b: _K) -> _V: ...
 def indexOf(__a: Iterable[_T], __b: _T) -> int: ...
 @overload
 def setitem(__a: MutableSequence[_T], __b: SupportsIndex, __c: _T) -> None: ...
@@ -106,17 +105,30 @@ class attrgetter(Generic[_T_co]):
 
 @final
 class itemgetter(Generic[_T_co]):
+    # mypy lacks support for PEP 646 https://github.com/python/mypy/issues/12280
+    # So we have to define all of these overloads to simulate unpacking the arguments
     @overload
-    def __new__(cls, item: Any) -> itemgetter[Any]: ...
+    def __new__(cls, item: _T_co) -> itemgetter[_T_co]: ...
     @overload
-    def __new__(cls, item: Any, __item2: Any) -> itemgetter[tuple[Any, Any]]: ...
+    def __new__(cls, item: _T_co, __item2: _T_co) -> itemgetter[tuple[_T_co, _T_co]]: ...
     @overload
-    def __new__(cls, item: Any, __item2: Any, __item3: Any) -> itemgetter[tuple[Any, Any, Any]]: ...
+    def __new__(cls, item: _T_co, __item2: _T_co, __item3: _T_co) -> itemgetter[tuple[_T_co, _T_co, _T_co]]: ...
     @overload
-    def __new__(cls, item: Any, __item2: Any, __item3: Any, __item4: Any) -> itemgetter[tuple[Any, Any, Any, Any]]: ...
+    def __new__(
+        cls, item: _T_co, __item2: _T_co, __item3: _T_co, __item4: _T_co
+    ) -> itemgetter[tuple[_T_co, _T_co, _T_co, _T_co]]: ...
     @overload
-    def __new__(cls, item: Any, *items: Any) -> itemgetter[tuple[Any, ...]]: ...
-    def __call__(self, obj: Any) -> _T_co: ...
+    def __new__(
+        cls, item: _T_co, __item2: _T_co, __item3: _T_co, __item4: _T_co, *items: _T_co
+    ) -> itemgetter[tuple[_T_co, ...]]: ...
+    # __key: _KT_contra in SupportsGetItem seems to be causing variance issues, ie:
+    # TypeVar "_KT_contra@SupportsGetItem" is contravariant
+    #   "tuple[int, int]" is incompatible with protocol "SupportsIndex"
+    # preventing [_T_co, ...] instead of [Any, ...]
+    #
+    # A suspected mypy issue prevents using [..., _T] instead of [..., Any] here.
+    # https://github.com/python/mypy/issues/14032
+    def __call__(self, obj: SupportsGetItem[Any, Any]) -> Any: ...
 
 @final
 class methodcaller:

stdlib/_typeshed/__init__.pyi

@@ -119,7 +119,7 @@ class SupportsKeysAndGetItem(Protocol[_KT, _VT_co]):
 
 # stable
 class SupportsGetItem(Protocol[_KT_contra, _VT_co]):
-    def __contains__(self, __x: object) -> bool: ...
+    def __contains__(self, __x: Any) -> bool: ...
     def __getitem__(self, __key: _KT_contra) -> _VT_co: ...
 
 # stable

stdlib/mmap.pyi

@@ -67,8 +67,11 @@ class mmap(Iterable[int], Sized):
     def __setitem__(self, __index: int, __object: int) -> None: ...
     @overload
     def __setitem__(self, __index: slice, __object: ReadableBuffer) -> None: ...
-    # Doesn't actually exist, but the object is actually iterable because it has __getitem__ and
-    # __len__, so we claim that there is also an __iter__ to help type checkers.
+    # Doesn't actually exist, but the object actually supports "in" because it has __getitem__,
+    # so we claim that there is also a __contains__ to help type checkers.
+    def __contains__(self, __o: object) -> bool: ...
+    # Doesn't actually exist, but the object is actually iterable because it has __getitem__ and __len__,
+    # so we claim that there is also an __iter__ to help type checkers.
     def __iter__(self) -> Iterator[int]: ...
     def __enter__(self: Self) -> Self: ...
     def __exit__(self, *args: object) -> None: ...

stubs/six/@tests/stubtest_allowlist.txt

@@ -8,8 +8,6 @@ six.create_bound_method.__defaults__
 six.moves.*
 
 # Implemented using "operator" functions in the implementation
-six.byte2int
-six.indexbytes
 six.get_function_closure
 six.get_function_code
 six.get_function_defaults
@@ -19,6 +17,8 @@ six.get_method_self
 six.viewitems
 six.viewkeys
 six.viewvalues
+# Should be `operator.itemgetter[int]`. But a bug in mypy prevents using TypeVar in itemgetter__call__
+six.byte2int
 
 # Unclear problems
 six.callable

stubs/six/six/__init__.pyi

@@ -1,7 +1,8 @@
 import builtins
+import operator
 import types
 import unittest
-from _typeshed import IdentityFunction
+from _typeshed import IdentityFunction, SupportsGetItem
 from builtins import next as next
 from collections.abc import Callable, ItemsView, Iterable, Iterator as _Iterator, KeysView, Mapping, ValuesView
 from functools import wraps as wraps
@@ -11,7 +12,7 @@ from re import Pattern
 from typing import Any, AnyStr, NoReturn, TypeVar, overload
 from typing_extensions import Literal
 
-from . import moves as moves
+from six import moves as moves
 
 _T = TypeVar("_T")
 _K = TypeVar("_K")
@@ -63,9 +64,13 @@ def u(s: str) -> str: ...
 unichr = chr
 
 def int2byte(i: int) -> bytes: ...
-def byte2int(bs: bytes) -> int: ...
-def indexbytes(buf: bytes, i: int) -> int: ...
-def iterbytes(buf: bytes) -> _Iterator[int]: ...
+
+# Should be `byte2int: operator.itemgetter[int]`. But a bug in mypy prevents using TypeVar in itemgetter__call__
+def byte2int(obj: SupportsGetItem[int, _T]) -> _T: ...
+
+indexbytes = operator.getitem
+iterbytes = iter
+
 def assertCountEqual(self: unittest.TestCase, first: Iterable[_T], second: Iterable[_T], msg: str | None = ...) -> None: ...
 @overload
 def assertRaisesRegex(self: unittest.TestCase, msg: str | None = ...) -> Any: ...

tests/stubtest_allowlists/py3_common.txt

@@ -615,6 +615,7 @@ wsgiref.handlers.BaseHandler.status
 # These would ideally be special-cased by type checkers; see https://github.com/python/mypy/issues/2220
 ctypes.Array.__iter__
 mmap.mmap.__iter__
+mmap.mmap.__contains__
 xml.etree.ElementTree.Element.__iter__
 xml.etree.cElementTree.Element.__iter__
 typing.IO.__iter__  # See https://github.com/python/typeshed/commit/97bc450acd60c1bcdafef3ce8fbe3b95a9c0cac3