refactoring after reviewers feedback + type hints

25-class-metaprog/persistent/dblib.py

@@ -3,17 +3,10 @@
 # Applying `check_identifier` to parameters prevents SQL injection.
 
 import sqlite3
-from typing import NamedTuple
+from typing import NamedTuple, Optional, Iterator, Any
 
 DEFAULT_DB_PATH = ':memory:'
-CONNECTION = None
-
-SQL_TYPES = {
-    int: 'INTEGER',
-    str: 'TEXT',
-    float: 'REAL',
-    bytes: 'BLOB',
-}
+CONNECTION: Optional[sqlite3.Connection] = None
 
 
 class NoConnection(Exception):
@@ -38,29 +31,45 @@ class UnexpectedMultipleResults(Exception):
     """Query returned more than 1 row."""
 
 
+SQLType = str
+
+TypeMap = dict[type, SQLType]
+
+SQL_TYPES: TypeMap = {
+    int: 'INTEGER',
+    str: 'TEXT',
+    float: 'REAL',
+    bytes: 'BLOB',
+}
+
+
 class ColumnSchema(NamedTuple):
     name: str
-    sql_type: str
+    sql_type: SQLType
 
 
-def check_identifier(name):
+FieldMap = dict[str, type]
+
+
+def check_identifier(name: str) -> None:
     if not name.isidentifier():
         raise ValueError(f'{name!r} is not an identifier')
 
 
-def connect(db_path=DEFAULT_DB_PATH):
+def connect(db_path: str = DEFAULT_DB_PATH) -> sqlite3.Connection:
     global CONNECTION
     CONNECTION = sqlite3.connect(db_path)
+    CONNECTION.row_factory = sqlite3.Row
     return CONNECTION
 
 
-def get_connection():
+def get_connection() -> sqlite3.Connection:
     if CONNECTION is None:
         raise NoConnection()
     return CONNECTION
 
 
-def gen_columns_sql(fields):
+def gen_columns_sql(fields: FieldMap) -> Iterator[ColumnSchema]:
     for name, py_type in fields.items():
         check_identifier(name)
         try:
@@ -70,7 +79,7 @@ def gen_columns_sql(fields):
         yield ColumnSchema(name, sql_type)
 
 
-def make_schema_sql(table_name, fields):
+def make_schema_sql(table_name: str, fields: FieldMap) -> str:
     check_identifier(table_name)
     pk = 'pk INTEGER PRIMARY KEY,'
     spcs = ' ' * 4
@@ -81,25 +90,24 @@ def make_schema_sql(table_name, fields):
     return f'CREATE TABLE {table_name} (\n{spcs}{pk}\n{spcs}{columns}\n)'
 
 
-def create_table(table_name, fields):
+def create_table(table_name: str, fields: FieldMap) -> None:
     con = get_connection()
     con.execute(make_schema_sql(table_name, fields))
 
 
-def read_columns_sql(table_name):
-    con = get_connection()
+def read_columns_sql(table_name: str) -> list[ColumnSchema]:
     check_identifier(table_name)
+    con = get_connection()
     rows = con.execute(f'PRAGMA table_info({table_name!r})')
-    # row fields: cid name type notnull dflt_value pk
-    return [ColumnSchema(r[1], r[2]) for r in rows]
+    return [ColumnSchema(r['name'], r['type']) for r in rows]
 
 
-def valid_table(table_name, fields):
+def valid_table(table_name: str, fields: FieldMap) -> bool:
     table_columns = read_columns_sql(table_name)
     return set(gen_columns_sql(fields)) <= set(table_columns)
 
 
-def ensure_table(table_name, fields):
+def ensure_table(table_name: str, fields: FieldMap) -> None:
     table_columns = read_columns_sql(table_name)
     if len(table_columns) == 0:
         create_table(table_name, fields)
@@ -107,21 +115,21 @@ def ensure_table(table_name, fields):
         raise SchemaMismatch(table_name)
 
 
-def insert_record(table_name, fields):
-    con = get_connection()
+def insert_record(table_name: str, data: dict[str, Any]) -> int:
     check_identifier(table_name)
-    placeholders = ', '.join(['?'] * len(fields))
+    con = get_connection()
+    placeholders = ', '.join(['?'] * len(data))
     sql = f'INSERT INTO {table_name} VALUES (NULL, {placeholders})'
-    cursor = con.execute(sql, tuple(fields.values()))
+    cursor = con.execute(sql, tuple(data.values()))
     pk = cursor.lastrowid
     con.commit()
     cursor.close()
     return pk
 
 
-def fetch_record(table_name, pk):
-    con = get_connection()
+def fetch_record(table_name: str, pk: int) -> sqlite3.Row:
     check_identifier(table_name)
+    con = get_connection()
     sql = f'SELECT * FROM {table_name} WHERE pk = ? LIMIT 2'
     result = list(con.execute(sql, (pk,)))
     if len(result) == 0:
@@ -132,19 +140,21 @@ def fetch_record(table_name, pk):
         raise UnexpectedMultipleResults()
 
 
-def update_record(table_name, pk, fields):
+def update_record(
+    table_name: str, pk: int, data: dict[str, Any]
+) -> tuple[str, tuple[Any, ...]]:
     check_identifier(table_name)
     con = get_connection()
-    names = ', '.join(fields.keys())
-    placeholders = ', '.join(['?'] * len(fields))
-    values = tuple(fields.values()) + (pk,)
+    names = ', '.join(data.keys())
+    placeholders = ', '.join(['?'] * len(data))
+    values = tuple(data.values()) + (pk,)
     sql = f'UPDATE {table_name} SET ({names}) = ({placeholders}) WHERE pk = ?'
     con.execute(sql, values)
     con.commit()
     return sql, values
 
 
-def delete_record(table_name, pk):
+def delete_record(table_name: str, pk: int) -> sqlite3.Cursor:
     con = get_connection()
     check_identifier(table_name)
     sql = f'DELETE FROM {table_name} WHERE pk = ?'