Add recipe path / weight.

Author: plypaul

metricflow-semantics/metricflow_semantics/experimental/semantic_graph/attribute_resolution/recipe_writer_path.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from typing_extensions import override
+
+from metricflow_semantics.experimental.mf_graph.path_finding.graph_path import MutableGraphPath
+from metricflow_semantics.experimental.mf_graph.path_finding.pathfinder import MetricflowPathfinder
+from metricflow_semantics.experimental.semantic_graph.attribute_resolution.attribute_recipe import (
+    AttributeRecipe,
+)
+from metricflow_semantics.experimental.semantic_graph.attribute_resolution.attribute_recipe_step import (
+    AttributeRecipeStep,
+)
+from metricflow_semantics.experimental.semantic_graph.sg_interfaces import SemanticGraphEdge, SemanticGraphNode
+from metricflow_semantics.mf_logging.pretty_formattable import MetricFlowPrettyFormattable
+from metricflow_semantics.mf_logging.pretty_formatter import PrettyFormatContext
+
+logger = logging.getLogger(__name__)
+
+_EMPTY_RECIPE = AttributeRecipe()
+
+
+@dataclass
+class AttributeRecipeWriterPath(MutableGraphPath[SemanticGraphNode, SemanticGraphEdge], MetricFlowPrettyFormattable):
+    """An implementation of a path in the semantic graph that writes recipes.
+
+    The nodes and edges in the semantic graph are annotated with recipe steps that describe the computation of
+    attributes. This path takes those steps and merges them into a single recipe as nodes are added to the path.
+
+    In DFS traversal, nodes are added and popped at the end. To support this operation, recipe versions are
+    stored in a list so that consistent state can be maintained.
+    """
+
+    # Every time a node / edge is added, the updated recipe is added to this list.
+    _recipe_versions: list[AttributeRecipe]
+
+    @staticmethod
+    def create(start_node: Optional[SemanticGraphNode] = None) -> AttributeRecipeWriterPath:  # noqa: D102
+        path = AttributeRecipeWriterPath(
+            _nodes=[],
+            _edges=[],
+            _weight_addition_order=[],
+            _current_weight=0,
+            _current_node_set=set(),
+            _node_set_addition_order=[],
+            _recipe_versions=[_EMPTY_RECIPE],
+        )
+        if start_node:
+            path._append_node(start_node)
+            path._append_step(start_node.recipe_step_to_append)
+        return path
+
+    @staticmethod
+    def create_from_edge(start_edge: SemanticGraphEdge, weight: int) -> AttributeRecipeWriterPath:  # noqa: D102
+        path = AttributeRecipeWriterPath(
+            _nodes=[],
+            _edges=[],
+            _weight_addition_order=[],
+            _current_weight=0,
+            _current_node_set=set(),
+            _node_set_addition_order=[],
+            _recipe_versions=[_EMPTY_RECIPE],
+        )
+        path.append_edge(start_edge, weight)
+        return path
+
+    def _append_step(self, recipe_step: AttributeRecipeStep) -> None:
+        previous_recipe = self._recipe_versions[-1]
+        self._recipe_versions.append(previous_recipe.append_step(recipe_step))
+
+    @property
+    def latest_recipe(self) -> AttributeRecipe:  # noqa: D102
+        return self._recipe_versions[-1]
+
+    @override
+    def append_edge(self, edge: SemanticGraphEdge, weight: int) -> None:
+        """Add an edge with the given weight to this path."""
+        if self.is_empty:
+            self._append_step(edge.tail_node.recipe_step_to_append)
+        self._append_step(edge.recipe_step_to_append)
+        self._append_step(edge.head_node.recipe_step_to_append)
+
+        super().append_edge(edge, weight)
+
+    @override
+    def pretty_format(self, format_context: PrettyFormatContext) -> Optional[str]:
+        return format_context.formatter.pretty_format(self._nodes)
+
+    @override
+    def pop_end(self) -> None:
+        if self._edges:
+            self._recipe_versions.pop()
+            self._recipe_versions.pop()
+        else:
+            self._recipe_versions.pop()
+        super().pop_end()
+
+    @override
+    def copy(self) -> AttributeRecipeWriterPath:
+        return AttributeRecipeWriterPath(
+            _nodes=self._nodes.copy(),
+            _edges=self._edges.copy(),
+            _current_weight=self._current_weight,
+            _current_node_set=self._current_node_set.copy(),
+            _weight_addition_order=self._weight_addition_order.copy(),
+            _node_set_addition_order=self._node_set_addition_order.copy(),
+            _recipe_versions=self._recipe_versions.copy(),
+        )
+
+
+RecipeWriterPathfinder = MetricflowPathfinder[SemanticGraphNode, SemanticGraphEdge, AttributeRecipeWriterPath]