Source code for prov.graph

from __future__ import annotations

import warnings
from typing import Any

import networkx as nx

from prov.identifier import QualifiedName
from prov.model import (
    PROV_ATTR_ACTIVITY,
    PROV_ATTR_AGENT,
    PROV_ATTR_ALTERNATE1,
    PROV_ATTR_ALTERNATE2,
    PROV_ATTR_BUNDLE,
    PROV_ATTR_COLLECTION,
    PROV_ATTR_DELEGATE,
    PROV_ATTR_ENDER,
    PROV_ATTR_ENTITY,
    PROV_ATTR_GENERAL_ENTITY,
    PROV_ATTR_GENERATED_ENTITY,
    PROV_ATTR_INFORMANT,
    PROV_ATTR_INFORMED,
    PROV_ATTR_PLAN,
    PROV_ATTR_RESPONSIBLE,
    PROV_ATTR_SPECIFIC_ENTITY,
    PROV_ATTR_STARTER,
    PROV_ATTR_TRIGGER,
    PROV_ATTR_USED_ENTITY,
    ProvActivity,
    ProvAgent,
    ProvBundle,
    ProvDocument,
    ProvElement,
    ProvEntity,
    ProvRecord,
    ProvRelation,
)

__author__ = "Trung Dong Huynh"
__email__ = "trungdong@donggiang.com"

warnings.warn(
    "In prov 3.0, graph export (prov.graph) will require the optional "
    '"graph" extra; install "prov[graph]" to keep using it after upgrading. '
    "See https://github.com/trungdong/prov/blob/master/ROADMAP.md",
    DeprecationWarning,
    stacklevel=2,
)


INFERRED_ELEMENT_CLASS = {
    PROV_ATTR_ENTITY: ProvEntity,
    PROV_ATTR_ACTIVITY: ProvActivity,
    PROV_ATTR_AGENT: ProvAgent,
    PROV_ATTR_TRIGGER: ProvEntity,
    PROV_ATTR_GENERATED_ENTITY: ProvEntity,
    PROV_ATTR_USED_ENTITY: ProvEntity,
    PROV_ATTR_DELEGATE: ProvAgent,
    PROV_ATTR_RESPONSIBLE: ProvAgent,
    PROV_ATTR_SPECIFIC_ENTITY: ProvEntity,
    PROV_ATTR_GENERAL_ENTITY: ProvEntity,
    PROV_ATTR_ALTERNATE1: ProvEntity,
    PROV_ATTR_ALTERNATE2: ProvEntity,
    PROV_ATTR_COLLECTION: ProvEntity,
    PROV_ATTR_INFORMED: ProvActivity,
    PROV_ATTR_INFORMANT: ProvActivity,
    PROV_ATTR_BUNDLE: ProvBundle,
    PROV_ATTR_PLAN: ProvEntity,
    PROV_ATTR_ENDER: ProvEntity,
    PROV_ATTR_STARTER: ProvEntity,
}
"""Maps a relation's first/second formal attribute (e.g. ``PROV_ATTR_ENTITY``)
to the :class:`~prov.model.ProvElement` subclass used to create a bare node
for an element referenced by a relation but not otherwise declared."""


[docs] def prov_to_graph(prov_document: ProvDocument) -> nx.MultiDiGraph[Any]: """Convert a :class:`~prov.model.ProvDocument` to a `MultiDiGraph <https://networkx.github.io/documentation/stable/reference/classes/multidigraph.html>`_ instance of the `NetworkX <https://networkx.github.io/>`_ library. The document is first :meth:`~prov.model.ProvBundle.unified` so that records referring to the same real-world thing are merged before the graph is built. Every PROV element becomes a graph node. Each relation becomes an edge between the elements named in its first two formal attributes, with the relation record stored under the edge's ``"relation"`` attribute; if either of those two elements was not already added as a node, a bare node is created for it using :data:`INFERRED_ELEMENT_CLASS`. Relations whose first two formal attributes are not both populated, or whose attribute type is not in :data:`INFERRED_ELEMENT_CLASS`, are silently skipped. Args: prov_document: The :class:`~prov.model.ProvDocument` instance to convert. Returns: A NetworkX ``MultiDiGraph`` with PROV elements as nodes and PROV relations as edges. """ g: nx.MultiDiGraph[Any] = nx.MultiDiGraph() unified = prov_document.unified() node_map: dict[QualifiedName | None, ProvRecord] = {} for element in unified.get_records(ProvElement): g.add_node(element) node_map[element.identifier] = element for relation in unified.get_records(ProvRelation): # taking the first two elements of a relation attr_pair_1, attr_pair_2 = relation.formal_attributes[:2] # only need the QualifiedName (i.e. the value of the attribute) qn1, qn2 = attr_pair_1[1], attr_pair_2[1] if qn1 and qn2: # only proceed if both ends of the relation exist try: if qn1 not in node_map: node_map[qn1] = INFERRED_ELEMENT_CLASS[attr_pair_1[0]](None, qn1) if qn2 not in node_map: node_map[qn2] = INFERRED_ELEMENT_CLASS[attr_pair_2[0]](None, qn2) except KeyError: # Unsupported attribute; cannot infer the type of the element continue # skipping this relation g.add_edge(node_map[qn1], node_map[qn2], relation=relation) return g
[docs] def graph_to_prov(g: nx.MultiDiGraph[Any]) -> ProvDocument: """Convert a `MultiDiGraph <https://networkx.github.io/documentation/stable/reference/classes/multidigraph.html>`_ that was previously produced by :func:`prov_to_graph` back to a :class:`~prov.model.ProvDocument`. Every node that is a :class:`~prov.model.ProvRecord` already attached to a bundle is added to the new document; nodes without a bundle (e.g. bare nodes inferred by :func:`prov_to_graph`) are skipped. Every edge whose ``"relation"`` data is a :class:`~prov.model.ProvRecord` is likewise added; edges without relation data are skipped. Args: g: The graph instance to convert. Returns: A new :class:`~prov.model.ProvDocument` containing the elements and relations recovered from ``g``. """ prov_doc = ProvDocument() for n in g.nodes(): if isinstance(n, ProvRecord) and n.bundle is not None: prov_doc.add_record(n) for _, _, edge_data in g.edges(data=True): try: relation = edge_data["relation"] if isinstance(relation, ProvRecord): prov_doc.add_record(relation) except KeyError: pass return prov_doc