stack/dev/scripts/generate_zot_tables.py

"""Introspect the Zotero SQLite database and generate Pydantic table models.

Usage:
    uv run python generate_zot_tables.py [--db PATH] [--out DIR]

Defaults:
    --db  ../zotero/data/zotero.sqlite
    --out ../src/bib/table/
"""

from __future__ import annotations

import argparse
import re
import sqlite3
from pathlib import Path

# ── SQLite type → Python type annotation mapping ─────────────────────────────

SQLITE_TYPE_MAP: dict[str, str] = {
    "TEXT": "str",
    "INT": "int",
    "INTEGER": "int",
    "REAL": "float",
    "BLOB": "bytes",
    "BOOLEAN": "bool",
    "TIMESTAMP": "str",
    "NONE": "str",
    "": "str",
}


def python_type(sqlite_type: str) -> str:
    """Map a SQLite column type to a Python type annotation string."""
    upper = sqlite_type.upper().strip()
    if upper in SQLITE_TYPE_MAP:
        return SQLITE_TYPE_MAP[upper]
    # Handle VARCHAR(N), CHAR(N), etc.
    if re.match(r"(VAR)?CHAR", upper):
        return "str"
    return "str"  # fallback


# ── Name helpers ─────────────────────────────────────────────────────────────


def to_class_name(table: str) -> str:
    """Convert a table name to PascalCase class name.

    camelCase table names like 'itemAttachments' → 'ItemAttachments'.
    snake_case table names like 'fulltext_items' → 'FulltextItems'.
    """
    # Split on underscores first
    parts = table.split("_")
    result = []
    for part in parts:
        if not part:
            continue
        # Split camelCase: 'itemAttachments' → ['item', 'Attachments']
        tokens = re.sub(r"([a-z])([A-Z])", r"\1_\2", part).split("_")
        for token in tokens:
            result.append(token.capitalize())
    return "".join(result)


def sanitize_field(name: str) -> str:
    """Ensure a column name is a valid Python identifier."""
    if name.isidentifier():
        return name
    cleaned = name.replace(" ", "_").replace("-", "_")
    if cleaned[0:1].isdigit():
        cleaned = f"f_{cleaned}"
    return cleaned


# ── Introspection ────────────────────────────────────────────────────────────


def get_tables(con: sqlite3.Connection) -> list[str]:
    """Return all user tables, sorted alphabetically."""
    rows = con.execute(
        "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
    ).fetchall()
    return [r[0] for r in rows]


def get_columns(con: sqlite3.Connection, table: str) -> list[tuple[str, str, bool]]:
    """Return [(col_name, sqlite_type, nullable), ...].

    SQLite PRAGMA table_info returns:
        cid, name, type, notnull, dflt_value, pk
    A column is nullable if notnull == 0 and it's not a primary key.
    """
    rows = con.execute(f'PRAGMA table_info("{table}")').fetchall()
    result = []
    for row in rows:
        col_name = row[1]
        col_type = row[2]
        notnull = row[3]
        is_pk = row[5]
        nullable = notnull == 0 and is_pk == 0
        result.append((col_name, col_type, nullable))
    return result


# ── Code generation ──────────────────────────────────────────────────────────


def generate_model(
    table: str, columns: list[tuple[str, str, bool]]
) -> tuple[str, set[str]]:
    """Generate a single Pydantic model class for a Zotero table."""
    cls = to_class_name(table)
    all_imports: set[str] = set()
    field_lines: list[str] = []

    for col_name, col_type, nullable in columns:
        py = python_type(col_type)
        fname = sanitize_field(col_name)
        if nullable:
            annotation = f"{py} | None = None"
        else:
            annotation = py
        field_lines.append(f"    {fname}: {annotation}")

    lines = [
        f"class {cls}(SQLTable):",
        f'    """Zotero table: {table}"""',
        "",
        '    __schema__ = "zotero"',
        f'    __tablename__ = "{table}"',
        "",
        *field_lines,
    ]

    return "\n".join(lines), all_imports


def generate_module(
    con: sqlite3.Connection,
    tables: list[str],
) -> str:
    """Generate a full Python module for all Zotero tables."""
    all_imports: set[str] = {
        "from __future__ import annotations",
        "from aco.table.base import SQLTable",
    }
    models: list[str] = []

    for table in tables:
        columns = get_columns(con, table)
        if not columns:
            continue
        model_code, imp = generate_model(table, columns)
        all_imports |= imp
        models.append(model_code)

    sorted_imports = sorted(all_imports)
    header = "\n".join(sorted_imports)
    body = "\n\n\n".join(models)

    return f"{header}\n\n\n{body}\n"


# ── Main ─────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate Pydantic models from Zotero SQLite"
    )
    from conf import path as _conf_path

    parser.add_argument(
        "--db",
        default=str(_conf_path("db.zotero")),
        help="Path to Zotero SQLite file",
    )
    parser.add_argument(
        "--out",
        default="../src/bib/table",
        help="Output directory",
    )
    args = parser.parse_args()

    db_path = Path(args.db)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    if not db_path.exists():
        print(f"Error: database not found at {db_path}")
        raise SystemExit(1)

    con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)

    tables = get_tables(con)
    print(f"Found {len(tables)} tables in {db_path}")

    # Generate single module with all tables
    module_path = out_dir / "zotero.py"
    code = generate_module(con, tables)
    module_path.write_text(code)
    print(f"Wrote {module_path} ({len(tables)} tables)")

    # Write __init__.py
    init_lines = [
        '"""Pydantic models for all Zotero SQLite tables."""',
        "",
        "from .zotero import *  # noqa: F401,F403",
        "",
    ]
    (out_dir / "__init__.py").write_text("\n".join(init_lines))
    print(f"Wrote {out_dir / '__init__.py'}")

    con.close()
    print(f"\nDone: {len(tables)} tables → {out_dir}/")


if __name__ == "__main__":
    main()