From 1707d8f466682bf89f9afc294123ec3b222d8e5f Mon Sep 17 00:00:00 2001
From: Daniel Bowman <daniel@danielbowman.co.uk>
Date: Sat, 17 Feb 2018 01:23:17 +0000
Subject: [PATCH] Improve pymysql types (#1804)

Add all cursor types and missing show_warnings method in connections

Added show_warnings to connections. No return type specified as this
requires tracing back through missing return types in previous
author's work. This is out of scope for this pull request.

Added SSDictCursor, SSCursor and DictCursorMixin to complete the
Cursors definition.
---
 third_party/2and3/pymysql/connections.pyi |  5 +++--
 third_party/2and3/pymysql/cursors.pyi     | 21 +++++++++++++++++----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/third_party/2and3/pymysql/connections.pyi b/third_party/2and3/pymysql/connections.pyi
index 098c17708..3789924b8 100644
--- a/third_party/2and3/pymysql/connections.pyi
+++ b/third_party/2and3/pymysql/connections.pyi
@@ -100,10 +100,10 @@ class Connection:
     def __enter__(self): ...
     def __exit__(self, exc, value, traceback): ...
     def query(self, sql): ...
-    def next_result(self): ...
+    def next_result(self, unbuffered: bool = ...): ...
     def affected_rows(self): ...
     def kill(self, thread_id): ...
-    def ping(self, reconnect=True): ...
+    def ping(self, reconnect: bool = ...): ...
     def set_charset(self, charset): ...
     def read_packet(self, packet_type=...): ...
     def insert_id(self): ...
@@ -112,6 +112,7 @@ class Connection:
     def get_host_info(self): ...
     def get_proto_info(self): ...
     def get_server_info(self): ...
+    def show_warnings(self): ...
     Warning = ...  # type: Any
     Error = ...  # type: Any
     InterfaceError = ...  # type: Any
diff --git a/third_party/2and3/pymysql/cursors.pyi b/third_party/2and3/pymysql/cursors.pyi
index be8cb90ee..80f37c101 100644
--- a/third_party/2and3/pymysql/cursors.pyi
+++ b/third_party/2and3/pymysql/cursors.pyi
@@ -1,4 +1,4 @@
-from typing import Union, Tuple, Any, Dict, Optional, Text
+from typing import Union, Tuple, Any, Dict, Optional, Text, Iterator, List
 from .connections import Connection
 
 Gen = Union[Tuple[Any, ...], Dict[str, Any]]
@@ -22,12 +22,25 @@ class Cursor:
     def executemany(self, query: str, args) -> int: ...
     def callproc(self, procname, args=...): ...
     def fetchone(self) -> Optional[Gen]: ...
-    def fetchmany(self, size: Optional[int] = ...) -> Optional[Gen]: ...
+    def fetchmany(self, size: Optional[int] = ...) -> Union[Optional[Gen], List[Gen]]: ...
     def fetchall(self) -> Optional[Tuple[Gen, ...]]: ...
-    def scroll(self, value, mode=''): ...
+    def scroll(self, value: int, mode: str = ...): ...
     def __iter__(self): ...
 
 class DictCursor(Cursor):
     def fetchone(self) -> Optional[Dict[str, Any]]: ...
-    def fetchmany(self, size=None) -> Optional[Tuple[Dict[str, Any], ...]]: ...
+    def fetchmany(self, size: Optional[int] = ...) -> Optional[Tuple[Dict[str, Any], ...]]: ...
     def fetchall(self) -> Optional[Tuple[Dict[str, Any], ...]]: ...
+
+class DictCursorMixin:
+    dict_type = ...  # type: Any
+
+class SSCursor(Cursor):
+    # fetchall return type is incompatible with the supertype.
+    def fetchall(self) -> List[Gen]: ...  # type: ignore
+    def fetchall_unbuffered(self) -> Iterator[Tuple[Gen, ...]]: ...
+    def __iter__(self) -> Iterator[Tuple[Gen, ...]]: ...
+    def fetchmany(self, size: Optional[int] = ...) -> List[Gen]: ...
+    def scroll(self, value: int, mode: str = ...) -> None: ...
+
+class SSDictCursor(DictCursorMixin, SSCursor): ...