test annotations to make sure thehy do something.

twisted · Nov 29, 2022 · b1e0777 · b1e0777
1 parent 3b8c7cf
commit b1e0777
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/src/twisted/internet/defer.py b/src/twisted/internet/defer.py
@@ -554,11 +554,21 @@ def addCallbacks(
         #     the caller should treat the result as the new type, consistently.
         return cast(Deferred[_NextDeferredResultT], self)
 
+    @overload
+    def addCallback(
+        self,
+        callback: Callable[Concatenate[_DeferredResultT, _P], Failure],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> Deferred[_NextDeferredResultT]:
+        ...
+
     @overload
     def addCallback(
         self,
         callback: Callable[
-            Concatenate[_DeferredResultT, _P], Union[_NextDeferredResultT, Failure]
+            Concatenate[_DeferredResultT, _P],
+            Union[Failure, Deferred[_NextDeferredResultT]],
         ],
         *args: _P.args,
         **kwargs: _P.kwargs,
@@ -568,7 +578,9 @@ def addCallback(
     @overload
     def addCallback(
         self,
-        callback: Callable[Concatenate[_DeferredResultT, _P], Failure],
+        callback: Callable[
+            Concatenate[_DeferredResultT, _P], Union[Failure, _NextDeferredResultT]
+        ],
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> Deferred[_NextDeferredResultT]:

diff --git a/src/twisted/test/test_defer.py b/src/twisted/test/test_defer.py
@@ -662,6 +662,36 @@ def test_callbackReturnsDeferred(self) -> None:
             self.callbackResults
         )
 
+    def test_callbackMaybeReturnsFailure(self) -> None:
+        """
+        Callbacks passed to addCallback may return Failures.
+        """
+        d: Deferred[int] = Deferred()
+
+        shouldFail = False
+        rte = RuntimeError()
+
+        def maybeFail(result: int) -> Union[int, Failure]:
+            # testing the Union return annotation here
+            if shouldFail:
+                return Failure(rte)
+            else:
+                return result + 1
+
+        d.callback(6)
+        self.assertEqual(self.successResultOf(d.addCallback(maybeFail)), 7)
+        d = Deferred[int]()
+        shouldFail = True
+        d.callback(6)
+        self.assertIs(
+            self.failureResultOf(
+                # make sure mypy still considers us a Deferred[int] here
+                d.addCallback(maybeFail).addCallback(maybeFail),
+                RuntimeError,
+            ).value,
+            rte,
+        )
+
     def test_errbackReturnsDeferred(self) -> None:
         """
         Errbacks passed to L{Deferred.addErrback} can return Deferreds just as