Improve several stdlib setdefault methods (#9612)

Author: AlexWaygood

stdlib/collections/__init__.pyi

@@ -354,7 +354,7 @@ class OrderedDict(dict[_KT, _VT], Reversible[_KT], Generic[_KT, _VT]):
     def fromkeys(cls, iterable: Iterable[_T], value: _S) -> OrderedDict[_T, _S]: ...
     # Keep OrderedDict.setdefault in line with MutableMapping.setdefault, modulo positional-only differences.
     @overload
-    def setdefault(self: OrderedDict[_KT, _T | None], key: _KT) -> _T | None: ...
+    def setdefault(self: OrderedDict[_KT, _T | None], key: _KT, default: None = None) -> _T | None: ...
     @overload
     def setdefault(self, key: _KT, default: _VT) -> _VT: ...
 
@@ -404,7 +404,11 @@ class ChainMap(MutableMapping[_KT, _VT], Generic[_KT, _VT]):
     def __contains__(self, key: object) -> bool: ...
     def __missing__(self, key: _KT) -> _VT: ...  # undocumented
     def __bool__(self) -> bool: ...
-    def setdefault(self, key: _KT, default: _VT = ...) -> _VT: ...
+    # Keep ChainMap.setdefault in line with MutableMapping.setdefault, modulo positional-only differences.
+    @overload
+    def setdefault(self: ChainMap[_KT, _T | None], key: _KT, default: None = None) -> _T | None: ...
+    @overload
+    def setdefault(self, key: _KT, default: _VT) -> _VT: ...
     @overload
     def pop(self, key: _KT) -> _VT: ...
     @overload

stdlib/typing.pyi

@@ -602,9 +602,13 @@ class MutableMapping(Mapping[_KT, _VT], Generic[_KT, _VT]):
     def pop(self, __key: _KT, default: _VT | _T) -> _VT | _T: ...
     def popitem(self) -> tuple[_KT, _VT]: ...
     # This overload should be allowed only if the value type is compatible with None.
-    # Keep OrderedDict.setdefault in line with MutableMapping.setdefault, modulo positional-only differences.
+    #
+    # Keep the following methods in line with MutableMapping.setdefault, modulo positional-only differences:
+    # -- collections.OrderedDict.setdefault
+    # -- collections.ChainMap.setdefault
+    # -- weakref.WeakKeyDictionary.setdefault
     @overload
-    def setdefault(self: MutableMapping[_KT, _T | None], __key: _KT) -> _T | None: ...
+    def setdefault(self: MutableMapping[_KT, _T | None], __key: _KT, __default: None = None) -> _T | None: ...
     @overload
     def setdefault(self, __key: _KT, __default: _VT) -> _VT: ...
     # 'update' used to take a Union, but using overloading is better.

stdlib/weakref.pyi

@@ -70,7 +70,7 @@ class WeakValueDictionary(MutableMapping[_KT, _VT]):
     def items(self) -> Iterator[tuple[_KT, _VT]]: ...  # type: ignore[override]
     def itervaluerefs(self) -> Iterator[KeyedRef[_KT, _VT]]: ...
     def valuerefs(self) -> list[KeyedRef[_KT, _VT]]: ...
-    def setdefault(self, key: _KT, default: _VT = ...) -> _VT: ...
+    def setdefault(self, key: _KT, default: _VT) -> _VT: ...  # type: ignore[override]
     @overload
     def pop(self, key: _KT) -> _VT: ...
     @overload
@@ -109,7 +109,11 @@ class WeakKeyDictionary(MutableMapping[_KT, _VT]):
     def values(self) -> Iterator[_VT]: ...  # type: ignore[override]
     def items(self) -> Iterator[tuple[_KT, _VT]]: ...  # type: ignore[override]
     def keyrefs(self) -> list[ref[_KT]]: ...
-    def setdefault(self, key: _KT, default: _VT = ...) -> _VT: ...
+    # Keep WeakKeyDictionary.setdefault in line with MutableMapping.setdefault, modulo positional-only differences
+    @overload
+    def setdefault(self: WeakKeyDictionary[_KT, _VT | None], key: _KT, default: None = None) -> _VT: ...
+    @overload
+    def setdefault(self, key: _KT, default: _VT) -> _VT: ...
     @overload
     def pop(self, key: _KT) -> _VT: ...
     @overload

test_cases/stdlib/typing/check_MutableMapping.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Any, Union
+from typing_extensions import assert_type
+
+
+def check_setdefault_method() -> None:
+    d: dict[int, str] = {}
+    d2: dict[int, str | None] = {}
+    d3: dict[int, Any] = {}
+
+    d.setdefault(1)  # type: ignore
+    assert_type(d.setdefault(1, "x"), str)
+    assert_type(d2.setdefault(1), Union[str, None])
+    assert_type(d2.setdefault(1, None), Union[str, None])
+    assert_type(d2.setdefault(1, "x"), Union[str, None])
+    assert_type(d3.setdefault(1), Union[Any, None])
+    assert_type(d3.setdefault(1, "x"), Any)

tests/stubtest_allowlists/py3_common.txt

@@ -214,6 +214,7 @@ weakref.ReferenceType.*  # Alias for _weakref.ReferenceType, problems should be
 weakref.WeakKeyDictionary.get
 weakref.WeakKeyDictionary.update
 weakref.WeakValueDictionary.get
+weakref.WeakValueDictionary.setdefault  # has a default value for the "default" argument, but always errors out if no value is supplied for the parameter by the user
 weakref.ref.*  # Alias for _weakref.ReferenceType, problems should be fixed there
 webbrowser.UnixBrowser.remote_action  # always overridden in inheriting class
 webbrowser.UnixBrowser.remote_action_newtab  # always overridden in inheriting class