PROV-JSON

def dumps_prov(graph: KnowledgeGraph, *, indent: int = 2) -> str:
    """Serialise a knowledge graph to a PROV-JSON string.

    Args:
        graph: The graph to export.
        indent: Indentation passed to :func:`json.dumps`.

    Returns:
        The PROV-JSON document as a string.

    """
    return json.dumps(to_prov(graph), indent=indent, ensure_ascii=False)

def from_prov(document: dict[str, Any]) -> KnowledgeGraph:
    """Reconstruct a knowledge graph from a PROV-JSON document.

    Args:
        document: A PROV-JSON document as produced by :func:`to_prov`.

    Returns:
        The reconstructed graph. External endpoint stubs (global ids or the
        generic ``ck:external`` / ``ck:node`` placeholders) are not materialised
        as nodes.

    """
    graph = KnowledgeGraph()
    for bucket in ("entity", "activity", "agent"):
        for qn, attrs in document.get(bucket, {}).items():
            local = _local(qn)
            if is_global_id(local) or attrs.get("prov:type") in ("ck:external", "ck:node"):
                continue
            graph.add_node(_node_from_attrs(qn, attrs))

    for name, payload_map in document.items():
        if name in _RELATION_ROLES_REV:
            etype, s_role, o_role = _RELATION_ROLES_REV[name]
            for qn, payload in payload_map.items():
                graph.add_edge(
                    Edge(type=etype, source=_local(payload[s_role]), target=_local(payload[o_role]), id=_local(qn))
                )
        elif name == "wasInfluencedBy":
            for qn, payload in payload_map.items():
                graph.add_edge(
                    Edge(
                        type=payload["ck:predicate"],
                        source=_local(payload["prov:influencee"]),
                        target=_local(payload["prov:influencer"]),
                        id=_local(qn),
                    )
                )
    return graph

def loads_prov(text: str) -> KnowledgeGraph:
    """Deserialise a knowledge graph from a PROV-JSON string.

    Args:
        text: A PROV-JSON document as produced by :func:`dumps_prov`.

    Returns:
        The reconstructed graph.

    """
    return from_prov(json.loads(text))

def to_prov(graph: KnowledgeGraph) -> dict[str, Any]:
    """Convert a knowledge graph to a PROV-JSON document.

    Args:
        graph: The graph to export.

    Returns:
        A PROV-JSON document as a plain dictionary.

    """
    entity: dict[str, Any] = {}
    activity: dict[str, Any] = {}
    agent: dict[str, Any] = {}
    buckets = {_ACTIVITY: activity, _AGENT: agent}

    for node in graph.nodes.values():
        buckets.get(node.type, entity)[_qn(node.id)] = _node_attrs(node)

    def _ensure_entity(node_id: str) -> None:
        qn = _qn(node_id)
        if qn not in entity and qn not in activity and qn not in agent:
            entity[qn] = {"prov:type": "ck:external"} if is_global_id(node_id) else {"prov:type": "ck:node"}

    relations: dict[str, dict[str, Any]] = {
        name: {} for name in ("used", "wasGeneratedBy", "wasDerivedFrom", "wasAttributedTo", "wasInfluencedBy")
    }
    for edge in graph.edges.values():
        _ensure_entity(edge.source)
        _ensure_entity(edge.target)
        subject, obj = _qn(edge.source), _qn(edge.target)
        roles = _RELATION_ROLES.get(edge.type)
        if roles is not None:
            name, s_role, o_role = roles
            relations[name][_qn(edge.id)] = {s_role: subject, o_role: obj}
        else:
            relations["wasInfluencedBy"][_qn(edge.id)] = {
                "prov:influencee": subject,
                "prov:influencer": obj,
                "ck:predicate": edge.type,
            }

    document: dict[str, Any] = {"prefix": {"ck": CK_NAMESPACE}}
    for name, collection in (("entity", entity), ("activity", activity), ("agent", agent), *relations.items()):
        if collection:
            document[name] = collection
    return document