diff --git a/cognite/neat/_session/_show.py b/cognite/neat/_session/_show.py index 13bf2d04b..d5ac64a3c 100644 --- a/cognite/neat/_session/_show.py +++ b/cognite/neat/_session/_show.py @@ -279,7 +279,7 @@ def __call__(self) -> Any: def _generate_dm_provenance_di_graph_and_types(self) -> nx.DiGraph: di_graph = nx.DiGraph() - hex_colored_types = _generate_hex_color_per_type(["Agent", "Entity", "Activity"]) + hex_colored_types = _generate_hex_color_per_type(["Agent", "Entity", "Activity", "Export", "Pruned"]) for change in self._state.rule_store.provenance: source = self._shorten_id(change.source_entity.id_) @@ -313,6 +313,33 @@ def _generate_dm_provenance_di_graph_and_types(self) -> nx.DiGraph: di_graph.add_edge(source, agent, label="used", color="grey") di_graph.add_edge(agent, target, label="generated", color="grey") + for source_id, exports in self._state.rule_store.exports_by_source_entity_id.items(): + source_shorten = self._shorten_id(source_id) + for export in exports: + export_id = self._shorten_id(export.target_entity.id_) + di_graph.add_node( + export_id, + label=export_id, + type="Export", + title="Export", + color=hex_colored_types["Export"], + ) + di_graph.add_edge(source_shorten, export_id, label="exported", color="grey") + + for pruned_lists in self._state.rule_store.pruned_by_source_entity_id.values(): + for prune_path in pruned_lists: + for change in prune_path: + source = self._shorten_id(change.source_entity.id_) + target = self._shorten_id(change.target_entity.id_) + di_graph.add_node( + target, + label=target, + type="Pruned", + title="Pruned", + color=hex_colored_types["Pruned"], + ) + di_graph.add_edge(source, target, label="pruned", color="grey") + return di_graph @staticmethod diff --git a/cognite/neat/_store/_rules_store.py b/cognite/neat/_store/_rules_store.py index 736424bfa..f6e04cc29 100644 --- a/cognite/neat/_store/_rules_store.py +++ b/cognite/neat/_store/_rules_store.py @@ -39,8 +39,8 @@ class OutcomeEntity(Entity): class NeatRulesStore: def __init__(self) -> None: self.provenance = Provenance() - self.exports_by_target_entity_id: dict[rdflib.URIRef, list[Change]] = defaultdict(list) - self.pruned_by_target_entity_id: dict[rdflib.URIRef, list[Provenance]] = defaultdict(list) + self.exports_by_source_entity_id: dict[rdflib.URIRef, list[Change]] = defaultdict(list) + self.pruned_by_source_entity_id: dict[rdflib.URIRef, list[Provenance]] = defaultdict(list) self._iteration_by_id: dict[Hashable, int] = {} def calculate_provenance_hash(self, shorten: bool = True) -> str: @@ -124,7 +124,7 @@ def export(self, exporter: BaseExporter[T_VerifiedRules, T_Export]) -> T_Export: description=exporter.description, source_entity=source_entity, ) - self.exports_by_target_entity_id[source_entity.id_].append(change) + self.exports_by_source_entity_id[source_entity.id_].append(change) return result def export_to_file(self, exporter: BaseExporter, path: Path) -> None: @@ -164,7 +164,7 @@ def export_to_file(self, exporter: BaseExporter, path: Path) -> None: description=exporter.description, source_entity=source_entity, ) - self.exports_by_target_entity_id[source_entity.id_].append(change) + self.exports_by_source_entity_id[source_entity.id_].append(change) def export_to_cdf(self, exporter: CDFExporter, client: NeatClient, dry_run: bool) -> UploadResultList: last_change = self.provenance[-1] @@ -205,7 +205,7 @@ def export_to_cdf(self, exporter: CDFExporter, client: NeatClient, dry_run: bool description=exporter.description, source_entity=source_entity, ) - self.exports_by_target_entity_id[source_entity.id_].append(change) + self.exports_by_source_entity_id[source_entity.id_].append(change) return result def prune_until_compatible(self, transformer: RulesTransformer) -> list[Change]: @@ -231,7 +231,7 @@ def prune_until_compatible(self, transformer: RulesTransformer) -> list[Change]: return [] self.provenance = self.provenance[: -len(pruned_candidates)] pruned_candidates.reverse() - self.pruned_by_target_entity_id[self.provenance[-1].target_entity.id_].append(Provenance(pruned_candidates)) + self.pruned_by_source_entity_id[self.provenance[-1].target_entity.id_].append(Provenance(pruned_candidates)) return pruned_candidates def _export(self, action: Callable[[Any], Any], agent: Agent, description: str) -> Any: