From 6aecafc64f02f54857bfa59b4eaab013e09be745 Mon Sep 17 00:00:00 2001
From: Petter Friberg <petter_friberg@hotmail.com>
Date: Tue, 25 Oct 2022 10:47:09 +0200
Subject: [PATCH] Improve type of `TestCase.captureOnCommitCallbacks` (#1202)

Co-authored-by: Adam Johnson <me@adamj.eu>
---
 django-stubs/test/testcases.pyi | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/django-stubs/test/testcases.pyi b/django-stubs/test/testcases.pyi
index d42abf7..5459253 100644
--- a/django-stubs/test/testcases.pyi
+++ b/django-stubs/test/testcases.pyi
@@ -1,5 +1,6 @@
 import threading
 import unittest
+from contextlib import contextmanager
 from datetime import date
 from types import TracebackType
 from typing import (
@@ -7,6 +8,7 @@ from typing import (
     Callable,
     Collection,
     Dict,
+    Generator,
     Iterable,
     Iterator,
     List,
@@ -216,7 +218,11 @@ class TransactionTestCase(SimpleTestCase):
 class TestCase(TransactionTestCase):
     @classmethod
     def setUpTestData(cls) -> None: ...
-    def captureOnCommitCallbacks(cls, *, using: Optional[str] = ..., execute: bool = ...): ...
+    @classmethod
+    @contextmanager
+    def captureOnCommitCallbacks(
+        cls, *, using: str = ..., execute: bool = ...
+    ) -> Generator[List[Callable[[], Any]], None, None]: ...
 
 class CheckCondition:
     conditions: Sequence[Tuple[Callable, str]] = ...