Fix unittest.mock.patch and unittest.mock.patch.object when new_callable is not None (#14358)

stdlib/@tests/test_cases/check_unittest.py

@@ -5,9 +5,9 @@ from collections.abc import Iterator, Mapping
 from datetime import datetime, timedelta
 from decimal import Decimal
 from fractions import Fraction
-from typing import TypedDict
+from typing import TypedDict, Union
 from typing_extensions import assert_type
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 case = unittest.TestCase()
 
@@ -154,10 +154,17 @@ def f_explicit_new(i: int) -> str:
     return "asdf"
 
 
+@patch("sys.exit", new_callable=lambda: 42)
+def f_explicit_new_callable(i: int, new_callable_ret: int) -> str:
+    return "asdf"
+
+
 assert_type(f_default_new(1), str)
 f_default_new("a")  # Not an error due to ParamSpec limitations
 assert_type(f_explicit_new(1), str)
 f_explicit_new("a")  # type: ignore[arg-type]
+assert_type(f_explicit_new_callable(1), str)
+f_explicit_new_callable("a")  # Same as default new
 
 
 @patch("sys.exit", new=Mock())
@@ -171,3 +178,51 @@ class TestXYZ(unittest.TestCase):
 
 assert_type(TestXYZ.attr, int)
 assert_type(TestXYZ.method(), int)
+
+
+with patch("sys.exit") as default_new_enter:
+    assert_type(default_new_enter, Union[MagicMock, AsyncMock])
+
+with patch("sys.exit", new=42) as explicit_new_enter:
+    assert_type(explicit_new_enter, int)
+
+with patch("sys.exit", new_callable=lambda: 42) as explicit_new_callable_enter:
+    assert_type(explicit_new_callable_enter, int)
+
+
+###
+# Tests for mock.patch.object
+###
+
+
+@patch.object(Decimal, "exp")
+def obj_f_default_new(i: int, mock: MagicMock) -> str:
+    return "asdf"
+
+
+@patch.object(Decimal, "exp", new=42)
+def obj_f_explicit_new(i: int) -> str:
+    return "asdf"
+
+
+@patch.object(Decimal, "exp", new_callable=lambda: 42)
+def obj_f_explicit_new_callable(i: int, new_callable_ret: int) -> str:
+    return "asdf"
+
+
+assert_type(obj_f_default_new(1), str)
+obj_f_default_new("a")  # Not an error due to ParamSpec limitations
+assert_type(obj_f_explicit_new(1), str)
+obj_f_explicit_new("a")  # type: ignore[arg-type]
+assert_type(obj_f_explicit_new_callable(1), str)
+obj_f_explicit_new_callable("a")  # Same as default new
+
+
+with patch.object(Decimal, "exp") as obj_default_new_enter:
+    assert_type(obj_default_new_enter, Union[MagicMock, AsyncMock])
+
+with patch.object(Decimal, "exp", new=42) as obj_explicit_new_enter:
+    assert_type(obj_explicit_new_enter, int)
+
+with patch.object(Decimal, "exp", new_callable=lambda: 42) as obj_explicit_new_callable_enter:
+    assert_type(obj_explicit_new_callable_enter, int)

stdlib/unittest/mock.pyi

@@ -262,7 +262,7 @@ class _patch(Generic[_T]):
 # This class does not exist at runtime, it's a hack to make this work:
 #     @patch("foo")
 #     def bar(..., mock: MagicMock) -> None: ...
-class _patch_default_new(_patch[MagicMock | AsyncMock]):
+class _patch_pass_arg(_patch[_T]):
     @overload
     def __call__(self, func: _TT) -> _TT: ...
     # Can't use the following as ParamSpec is only allowed as last parameter:
@@ -303,7 +303,7 @@ class _patcher:
         create: bool = ...,
         spec_set: Any | None = ...,
         autospec: Any | None = ...,
-        new_callable: Any | None = ...,
+        new_callable: Callable[..., Any] | None = ...,
         **kwargs: Any,
     ) -> _patch[_T]: ...
     @overload
@@ -315,9 +315,21 @@ class _patcher:
         create: bool = ...,
         spec_set: Any | None = ...,
         autospec: Any | None = ...,
-        new_callable: Any | None = ...,
+        new_callable: Callable[..., _T],
         **kwargs: Any,
-    ) -> _patch_default_new: ...
+    ) -> _patch_pass_arg[_T]: ...
+    @overload
+    def __call__(
+        self,
+        target: str,
+        *,
+        spec: Any | None = ...,
+        create: bool = ...,
+        spec_set: Any | None = ...,
+        autospec: Any | None = ...,
+        new_callable: None = ...,
+        **kwargs: Any,
+    ) -> _patch_pass_arg[MagicMock | AsyncMock]: ...
     @overload
     @staticmethod
     def object(
@@ -328,7 +340,7 @@ class _patcher:
         create: bool = ...,
         spec_set: Any | None = ...,
         autospec: Any | None = ...,
-        new_callable: Any | None = ...,
+        new_callable: Callable[..., Any] | None = ...,
         **kwargs: Any,
     ) -> _patch[_T]: ...
     @overload
@@ -341,9 +353,22 @@ class _patcher:
         create: bool = ...,
         spec_set: Any | None = ...,
         autospec: Any | None = ...,
-        new_callable: Any | None = ...,
+        new_callable: Callable[..., _T],
         **kwargs: Any,
-    ) -> _patch[MagicMock | AsyncMock]: ...
+    ) -> _patch_pass_arg[_T]: ...
+    @overload
+    @staticmethod
+    def object(
+        target: Any,
+        attribute: str,
+        *,
+        spec: Any | None = ...,
+        create: bool = ...,
+        spec_set: Any | None = ...,
+        autospec: Any | None = ...,
+        new_callable: None = ...,
+        **kwargs: Any,
+    ) -> _patch_pass_arg[MagicMock | AsyncMock]: ...
     @staticmethod
     def multiple(
         target: Any,