From 58b26fdbd34a8f688bc236925c367df6b180e2a9 Mon Sep 17 00:00:00 2001
From: Ran Benita <ran@unusedvar.com>
Date: Tue, 10 Dec 2019 18:37:13 +0200
Subject: [PATCH] Improve TestCase.assertNumQueries type (#250)

---
 django-stubs/test/testcases.pyi | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/django-stubs/test/testcases.pyi b/django-stubs/test/testcases.pyi
index c7123a2..5d9b3ff 100644
--- a/django-stubs/test/testcases.pyi
+++ b/django-stubs/test/testcases.pyi
@@ -1,7 +1,7 @@
 import threading
 import unittest
 from datetime import date
-from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, ClassVar
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, ClassVar, overload
 
 from django.core.exceptions import ImproperlyConfigured
 from django.core.handlers.wsgi import WSGIHandler
@@ -146,9 +146,14 @@ class TransactionTestCase(SimpleTestCase):
         ordered: bool = ...,
         msg: Optional[str] = ...,
     ) -> None: ...
+    @overload
     def assertNumQueries(
-        self, num: int, func: Optional[Union[Callable, Type[list]]] = ..., *args: Any, using: Any = ..., **kwargs: Any
-    ) -> Optional[_AssertNumQueriesContext]: ...
+        self, num: int, func: Callable[..., Any], *args: Any, using: str = ..., **kwargs: Any
+    ) -> None: ...
+    @overload
+    def assertNumQueries(
+        self, num: int, func: None = ..., *args: Any, using: str = ..., **kwargs: Any
+    ) -> _AssertNumQueriesContext: ...
 
 class TestCase(TransactionTestCase):
     @classmethod