From 5b47bc956ec388847f503b540a6ed9d9480d734b Mon Sep 17 00:00:00 2001
From: Stefan Urbanek <stefan.urbanek@gmail.com>
Date: Tue, 21 Mar 2017 23:11:41 -0700
Subject: [PATCH] Added SQLAlchemy annotations (#1029)

* Added SQLAlchemy annotations

* Made Connection and Engine sublcasses of Connectable (python/typeshed#1018)
* Moved execute() from Connection to Connectable
* Made RowProxy a Mapping and removed Mapping inherited methods
* Made ResultProxy an Iterator of RowProxy
* Added most relevant methods for fetching of ResultProxy
* Added where(), group_by(), order_by() and limit() to Select

* Follow squalchemy module structure

* Created sqlalchemy.engine.result and moved ResultProxy and RowProxy
  there
* Created sqlalchemy.engine.interfaces and moved Connectable there
* Added non-deprecated methods to Connectable: connect,
  contextual_connect and scalar
* Fixed return type of scalar() to Any

* Missed ResultProxy scalar return

... had it in Connectable only.
---
 .../2and3/sqlalchemy/engine/__init__.pyi      |  2 +-
 third_party/2and3/sqlalchemy/engine/base.pyi  | 17 +++-----------
 .../2and3/sqlalchemy/engine/interfaces.pyi    | 13 +++++++++++
 .../2and3/sqlalchemy/engine/result.pyi        | 22 +++++++++++++++++++
 .../2and3/sqlalchemy/sql/selectable.pyi       |  8 ++++++-
 5 files changed, 46 insertions(+), 16 deletions(-)
 create mode 100644 third_party/2and3/sqlalchemy/engine/interfaces.pyi
 create mode 100644 third_party/2and3/sqlalchemy/engine/result.pyi

diff --git a/third_party/2and3/sqlalchemy/engine/__init__.pyi b/third_party/2and3/sqlalchemy/engine/__init__.pyi
index 49eca6d67..396cac5ca 100644
--- a/third_party/2and3/sqlalchemy/engine/__init__.pyi
+++ b/third_party/2and3/sqlalchemy/engine/__init__.pyi
@@ -4,8 +4,8 @@
 
 from .base import Connection as Connection
 from .base import Engine as Engine
-from .base import RowProxy as RowProxy
 from .base import Transaction as Transaction
+from .result import RowProxy as RowProxy
 
 def create_engine(*args, **kwargs): ...
 def engine_from_config(configuration, prefix=..., **kwargs): ...
diff --git a/third_party/2and3/sqlalchemy/engine/base.pyi b/third_party/2and3/sqlalchemy/engine/base.pyi
index 71f6e196a..af3cf8c54 100644
--- a/third_party/2and3/sqlalchemy/engine/base.pyi
+++ b/third_party/2and3/sqlalchemy/engine/base.pyi
@@ -1,20 +1,9 @@
-from typing import Any, List, Tuple
+from .interfaces import Connectable
 
-# Dummy until I figure out something better.
-class Connectable:
-    pass
-
-class Connection:
+class Connection(Connectable):
     def begin(self): ...
-    def execute(self, object, *multiparams, **params): ...
 
-class Engine(object): ...
-
-class RowProxy:
-    def items(self) -> List[Tuple[Any, Any]]: ...
-    def keys(self) -> List[Any]: ...
-    def values(self) -> List[Any]: ...
-    def __getitem__(self, key: str): ...
+class Engine(Connectable): ...
 
 class Transaction:
     def commit(self): ...
diff --git a/third_party/2and3/sqlalchemy/engine/interfaces.pyi b/third_party/2and3/sqlalchemy/engine/interfaces.pyi
new file mode 100644
index 000000000..c35c05a8e
--- /dev/null
+++ b/third_party/2and3/sqlalchemy/engine/interfaces.pyi
@@ -0,0 +1,13 @@
+from typing import Any, TYPE_CHECKING
+from .result import ResultProxy
+
+if TYPE_CHECKING:
+    from .base import Connection
+
+class Connectable:
+    def execute(self, object, *multiparams: Any, **params: Any) -> ResultProxy: ...
+    def connect(self, **kwargs: Any) -> Connection: ...
+    def contextual_connect(self) -> Connection: ...
+    # Note: The return type `Any` should be a DB API 2 value type once defined
+    # TODO: See #1037
+    def scalar(self, object, *multiparams: Any, **params: Any) -> Any: ...
diff --git a/third_party/2and3/sqlalchemy/engine/result.pyi b/third_party/2and3/sqlalchemy/engine/result.pyi
new file mode 100644
index 000000000..0361ee05e
--- /dev/null
+++ b/third_party/2and3/sqlalchemy/engine/result.pyi
@@ -0,0 +1,22 @@
+from typing import Any, Iterator, List, Mapping, Optional
+
+# Note: The value type `Any` should be a DB API 2 value type once defined
+# TODO: See #1037
+
+class RowProxy(Mapping[str, Any]): ...
+
+class ResultProxy(Iterator[RowProxy]):
+    def keys(self) -> List[str]: ...
+    def close(self) -> None: ...
+    def __iter__(self) -> Iterator[RowProxy]: ...
+    def fetchall(self) -> Iterator[RowProxy]: ...
+    def fetchmany(self, size: Optional[int]=None) -> Iterator[RowProxy]: ...
+    def fetchone(self) -> Optional[RowProxy]: ...
+    def first(self) -> Optional[RowProxy]: ...
+
+    # Note: The return type `Any` should be a DB API 2 value type once defined
+    # TODO: See #1037
+    def scalar(self) -> Any: ...
+
+    @property
+    def rowcount(self) -> int: ...
diff --git a/third_party/2and3/sqlalchemy/sql/selectable.pyi b/third_party/2and3/sqlalchemy/sql/selectable.pyi
index 543349d64..b12feb47e 100644
--- a/third_party/2and3/sqlalchemy/sql/selectable.pyi
+++ b/third_party/2and3/sqlalchemy/sql/selectable.pyi
@@ -69,7 +69,13 @@ class ForUpdateArg(ClauseElement): ...
 class SelectBase(Executable, FromClause): ...
 class GenerativeSelect(SelectBase): ...
 class CompoundSelect(GenerativeSelect): ...
-class Select(HasPrefixes, HasSuffixes, GenerativeSelect): ...
+
+class Select(HasPrefixes, HasSuffixes, GenerativeSelect):
+    def where(self, whereclause) -> Select: ...
+    def group_by(self, *clauses) -> Select: ...
+    def order_by(self, *clauses) -> Select: ...
+    def limit(self, limit: int) -> Select: ...
+
 class ScalarSelect(Generative, Grouping): ...
 class Exists(UnaryExpression): ...
 class TextAsFrom(SelectBase): ...