From 3931432b34ac72db19c2e25ccdcb1a7fe598f65b Mon Sep 17 00:00:00 2001
From: LanDinh <coding+github@khaleesi.ninja>
Date: Sun, 18 Apr 2021 12:50:25 +0200
Subject: [PATCH] Add type hints for the postgres CursorDebugWrapper and add
 missing method to the BaseCursorDebugWrapper (#585)

* Add type hints for the postgres CursorDebugWrapper, expand the BaseCursorDebugWrapper.

* Fix how Optinal gets applied.

* Fix optional handling further.

* Adjust postgres debugcursorwrapper to look more like the implementation.

* Apply review feedback.

Co-authored-by: LanDinh <coding+sourcetree@khaleesi.ninja>
---
 django-stubs/db/backends/postgresql/base.pyi |  8 +++++++-
 django-stubs/db/backends/utils.pyi           | 18 +++++++++++-------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/django-stubs/db/backends/postgresql/base.pyi b/django-stubs/db/backends/postgresql/base.pyi
index 91674f3..8f4a757 100644
--- a/django-stubs/db/backends/postgresql/base.pyi
+++ b/django-stubs/db/backends/postgresql/base.pyi
@@ -1,6 +1,8 @@
-from typing import Dict, Tuple
+from io import IOBase
+from typing import Any, Dict, Optional, Tuple
 
 from django.db.backends.base.base import BaseDatabaseWrapper
+from django.db.backends.utils import CursorDebugWrapper as BaseCursorDebugWrapper
 
 def psycopg2_version() -> Tuple[int, ...]: ...
 
@@ -15,3 +17,7 @@ class DatabaseWrapper(BaseDatabaseWrapper):
     _named_cursor_idx: int = ...
     @property
     def pg_version(self) -> str: ...
+
+class CursorDebugWrapper(BaseCursorDebugWrapper):
+    def copy_expert(self, sql: str, file: IOBase, *args: Any): ...
+    def copy_to(self, file: IOBase, table: str, *args: Any, **kwargs: Any): ...
diff --git a/django-stubs/db/backends/utils.pyi b/django-stubs/db/backends/utils.pyi
index a867b51..4b01949 100644
--- a/django-stubs/db/backends/utils.pyi
+++ b/django-stubs/db/backends/utils.pyi
@@ -1,13 +1,14 @@
 import types
 from datetime import date, datetime, time
 from decimal import Decimal
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
+from typing import Any, ContextManager, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 from uuid import UUID
 
 logger: Any
 
 # Python types that can be adapted to SQL.
 _SQLType = Union[None, bool, int, float, Decimal, str, bytes, datetime, UUID]
+_ExecuteParameters = Optional[Union[Sequence[_SQLType], Mapping[str, _SQLType]]]
 
 class CursorWrapper:
     cursor: Any = ...
@@ -24,16 +25,19 @@ class CursorWrapper:
         tb: Optional[types.TracebackType],
     ) -> None: ...
     def callproc(self, procname: str, params: List[Any] = ..., kparams: Dict[str, int] = ...) -> Any: ...
-    def execute(
-        self, sql: str, params: Optional[Union[Sequence[_SQLType], Mapping[str, _SQLType]]] = ...
-    ) -> Optional[Any]: ...
-    def executemany(
-        self, sql: str, param_list: Sequence[Optional[Union[Sequence[_SQLType], Mapping[str, _SQLType]]]]
-    ) -> Optional[Any]: ...
+    def execute(self, sql: str, params: _ExecuteParameters = ...) -> Any: ...
+    def executemany(self, sql: str, param_list: Sequence[_ExecuteParameters]) -> Any: ...
 
 class CursorDebugWrapper(CursorWrapper):
     cursor: Any
     db: Any
+    def debug_sql(
+        self,
+        sql: Optional[str] = ...,
+        params: Optional[Union[_ExecuteParameters, Sequence[_ExecuteParameters]]] = ...,
+        use_last_executed_query: bool = ...,
+        many: bool = ...,
+    ) -> ContextManager[None]: ...
 
 def typecast_date(s: Optional[str]) -> Optional[date]: ...
 def typecast_time(s: Optional[str]) -> Optional[time]: ...