basetypes: better DictRow

Author: markokr

skytools/basetypes.py

@@ -2,27 +2,43 @@
 """
 
 import io
-from typing import Sequence, Mapping, List, Any, Optional, Union, IO
-
+from typing import IO, Any, Iterable, Mapping, Optional, Sequence, Tuple, Union
 
 try:
     from typing import Protocol
 except ImportError:
     Protocol = object   # type: ignore
 
-RowType = Sequence[Any]
+
+ExecuteParams = Union[Sequence[Any], Mapping[str, Any]]
+
+
+class DictRow(Protocol):
+    """Allow both key and index-based access.
+
+    Both Psycopg2 DictRow and PL/Python rows support this.
+    """
+    def keys(self) -> Iterable[str]: ...
+    def values(self) -> Iterable[Any]: ...
+    def items(self) -> Iterable[Tuple[str, Any]]: ...
+    def __getitem__(self, key: Union[str, int]) -> Any: ...
+    def __iter__(self) -> Iterable[str]: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, key: str) -> bool: ...
 
 
 class Cursor(Protocol):
-    def execute(self, sql: str, params: Optional[Union[Sequence[Any], Mapping[str, Any]]] = None) -> None: ...
-    def fetchall(self) -> List[RowType]: ...
-    def fetchone(self) -> RowType: ...
+    def execute(self, sql: str, params: Optional[ExecuteParams] = None) -> None: ...
+    def fetchall(self) -> Sequence[DictRow]: ...
+    def fetchone(self) -> DictRow: ...
     def copy_from(self, buf: IO[str], hdr: str) -> None: ...
     def copy_expert(self, sql: str, f: Union[IO[str], io.TextIOBase]) -> None: ...
 
 
 class Connection(Protocol):
     def cursor(self) -> Cursor: ...
+    def rollback(self) -> None: ...
+    def commit(self) -> None: ...
 
 
 class Runnable(Protocol):

skytools/querybuilder.py

@@ -16,7 +16,7 @@
 
 import skytools
 
-from .sqltools import Cursor
+from .basetypes import Cursor, DictRow
 
 try:
     import plpy
@@ -278,7 +278,7 @@ def __init__(self, sql: str):
         self.arg_map = qb._arg_value_list
         self.sql = sql
 
-    def execute(self, arg_dict: Mapping[str, Any], all_keys_required=True) -> List[Mapping[str, Any]]:
+    def execute(self, arg_dict: Mapping[str, Any], all_keys_required=True) -> Sequence[DictRow]:
         try:
             if all_keys_required:
                 arg_list = [arg_dict[k] for k in self.arg_map]
@@ -300,7 +300,7 @@ def __repr__(self) -> str:
 def plpy_exec(gd: Optional[Dict[str, Any]],
               sql: str,
               args: Optional[Mapping[str, Any]],
-              all_keys_required=True) -> Sequence[Mapping[str, Any]]:
+              all_keys_required=True) -> Sequence[DictRow]:
     """Cached plan execution for PL/Python.
 
     @param gd:  dict to store cached plans under.  If None, caching is disabled.