auth: Add some typing information to the functions

1 files changed, 42 insertions, 2 deletions

diff --git a/gn3/auth/db.py b/gn3/auth/db.py
index e732a03..8760153 100644
--- a/gn3/auth/db.py
+++ b/gn3/auth/db.py
@@ -1,9 +1,49 @@
 """Handle connection to auth database."""
 import sqlite3
 import contextlib
+from typing import Any, Iterator, Protocol
+
+class DbConnection(Protocol):
+    """Type annotation for a generic database connection object."""
+    def cursor(self) -> Any:
+        """A cursor object"""
+        ...
+
+    def commit(self) -> Any:
+        """Commit the transaction."""
+        ...
+
+    def rollback(self) -> Any:
+        """Rollback the transaction."""
+        ...
+
+class DbCursor(Protocol):
+    """Type annotation for a generic database cursor object."""
+    def execute(self, *args, **kwargs) -> Any:
+        """Execute a single query"""
+        ...
+
+    def executemany(self, *args, **kwargs) -> Any:
+        """
+        Execute parameterized SQL statement sql against all parameter sequences
+        or mappings found in the sequence parameters.
+        """
+        ...
+
+    def fetchone(self, *args, **kwargs):
+        """Fetch single result if present, or `None`."""
+        ...
+
+    def fetchmany(self, *args, **kwargs):
+        """Fetch many results if present or `None`."""
+        ...
+
+    def fetchall(self, *args, **kwargs):
+        """Fetch all results if present or `None`."""
+        ...
 
 @contextlib.contextmanager
-def connection(db_path: str):
+def connection(db_path: str) -> Iterator[DbConnection]:
     """Create the connection to the auth database."""
     conn = sqlite3.connect(db_path)
     try:
@@ -15,7 +55,7 @@ def connection(db_path: str):
         conn.close()
 
 @contextlib.contextmanager
-def cursor(conn):
+def cursor(conn: DbConnection) -> Iterator[DbCursor]:
     """Get a cursor from the given connection to the auth database."""
     cur = conn.cursor()
     try: