typing: Split assertRaises context manager overload

This ensures we don't lose type information on the `exception` attribute
of the context manager:

  with self.assertRaises(ValueError) as ctx:
      int("foo")

  reveal_type(ctx.exception)  # now reveals ValueError

We also narrow the tuple type, but varadic expansion is not a thing yet,
so we need explicitly type for different lengths of tuple. We decide
that at most 3 exceptions can be passed: any more and you'll lose this
information, which seems fair.

Signed-off-by: Stephen Finucane <stephen@that.guru>

Author: stephenfin

testtools/testcase.py

@@ -28,6 +28,9 @@
 
 T = TypeVar("T")
 U = TypeVar("U")
+_E = TypeVar("_E", bound=BaseException)
+_E2 = TypeVar("_E2", bound=BaseException)
+_E3 = TypeVar("_E3", bound=BaseException)
 
 # ruff: noqa: E402 - TypeVars must be defined before importing testtools modules
 from testtools import content
@@ -513,17 +516,38 @@ def assertRaises(
     @overload  # type: ignore[override]
     def assertRaises(
         self,
-        expected_exception: type[BaseException] | tuple[type[BaseException]],
+        expected_exception: type[_E],
         callable: None = ...,
-    ) -> "_AssertRaisesContext": ...
+    ) -> "_AssertRaisesContext[_E]": ...
 
-    def assertRaises(  # type: ignore[override]
+    @overload  # type: ignore[override]
+    def assertRaises(
         self,
-        expected_exception: type[BaseException] | tuple[type[BaseException]],
+        expected_exception: tuple[type[_E], type[_E2]],
+        callable: None = ...,
+    ) -> "_AssertRaisesContext[_E | _E2]": ...
+
+    @overload  # type: ignore[override]
+    def assertRaises(
+        self,
+        expected_exception: tuple[type[_E], type[_E2], type[_E3]],
+        callable: None = ...,
+    ) -> "_AssertRaisesContext[_E | _E2 | _E3]": ...
+
+    @overload  # type: ignore[override]
+    def assertRaises(
+        self,
+        expected_exception: tuple[type[BaseException], ...],
+        callable: None = ...,
+    ) -> "_AssertRaisesContext[BaseException]": ...
+
+    def assertRaises(  # type: ignore[override, misc]
+        self,
+        expected_exception: type[BaseException] | tuple[type[BaseException], ...],
         callable: Callable[_P, _R] | None = None,
         *args: _P.args,
         **kwargs: _P.kwargs,
-    ) -> "_AssertRaisesContext | BaseException":
+    ) -> "_AssertRaisesContext[BaseException] | BaseException":
         """Fail unless an exception of class expected_exception is thrown
         by callable when invoked with arguments args and keyword
         arguments kwargs. If a different type of exception is
@@ -1223,15 +1247,15 @@ def _id(obj: _F) -> _F:
     return _id
 
 
-class _AssertRaisesContext:
+class _AssertRaisesContext(Generic[_E]):
     """A context manager to handle expected exceptions for assertRaises.
 
     This provides compatibility with unittest's assertRaises context manager.
     """
 
     def __init__(
         self,
-        expected: type[BaseException] | tuple[type[BaseException]],
+        expected: type[_E] | tuple[type[BaseException], ...],
         test_case: TestCase,
         msg: str | None = None,
     ) -> None:
@@ -1244,7 +1268,7 @@ def __init__(
         self.expected = expected
         self.test_case = test_case
         self.msg = msg
-        self.exception: BaseException | None = None
+        self.exception: _E | None = None
 
     def __enter__(self) -> "Self":
         return self
@@ -1274,7 +1298,7 @@ def __exit__(
             # let unexpected exceptions pass through
             return False
         # store exception for later retrieval
-        self.exception = exc_value
+        self.exception = cast(_E, exc_value)
         return True