backend/hypergraph/persistence.py

import sqlite3, json
from ..config import HYPERGRAPH_DB
from .engine import Hypergraph


def save(hg: Hypergraph):
    con = sqlite3.connect(HYPERGRAPH_DB)
    con.executescript("""
        CREATE TABLE IF NOT EXISTS nodes(node_id TEXT PRIMARY KEY, data TEXT);
        CREATE TABLE IF NOT EXISTS edges(edge_id TEXT PRIMARY KEY, src TEXT, dst TEXT, data TEXT);
        DELETE FROM nodes; DELETE FROM edges;
    """)
    for i in hg.g.node_indexes():
        d = hg.g[i]
        con.execute("INSERT OR REPLACE INTO nodes VALUES(?,?)", (d["node_id"], json.dumps(d)))
    for s, t in hg.g.edge_list():
        d = hg.g.get_edge_data(s, t)
        sd, td = hg.g[s]["node_id"], hg.g[t]["node_id"]
        con.execute("INSERT OR REPLACE INTO edges VALUES(?,?,?,?)", (d["edge_id"], sd, td, json.dumps(d)))
    con.commit(); con.close()


def load() -> Hypergraph:
    hg = Hypergraph()
    if not HYPERGRAPH_DB.exists():
        return hg
    con = sqlite3.connect(HYPERGRAPH_DB)
    try:
        for nid, data in con.execute("SELECT node_id, data FROM nodes"):
            d = json.loads(data)
            i = hg.g.add_node(d)
            hg._idx[nid] = i
        for eid, sd, td, data in con.execute("SELECT edge_id, src, dst, data FROM edges"):
            d = json.loads(data)
            if sd in hg._idx and td in hg._idx:
                hg.g.add_edge(hg._idx[sd], hg._idx[td], d)
    except sqlite3.OperationalError:
        pass
    con.close()
    return hg