refactor: render spec line recursively on rust side

Author: lemorage

python/cocoindex/flow.py

@@ -458,38 +458,19 @@ def _render_spec(self, verbose: bool = False) -> Tree:
         """
         Render the flow spec as a styled rich Tree with hierarchical structure.
         """
+        spec = self._get_spec(verbose=verbose)
         tree = Tree(f"Flow: {self.name}", style="cyan")
-        current_section = None
-        section_node = None
-        indent_stack = []
-
-        for i, (section, content, indent) in enumerate(self._get_spec(verbose=verbose)):
-            # Skip "Scope" entries (see ReactiveOpScope in spec.rs)
-            if content.startswith("Scope:"):
-                continue
-
-            if section != current_section:
-                current_section = section
-                section_node = tree.add(f"{section}:", style="bold magenta")
-                indent_stack = [(0, section_node)]
-
-            while indent_stack and indent_stack[-1][0] >= indent:
-                indent_stack.pop()
-
-            parent = indent_stack[-1][1] if indent_stack else section_node
-            styled_content = Text(content, style="yellow")
-            is_parent = any(
-                next_indent > indent
-                for _, next_content, next_indent in self._get_spec(verbose=verbose)[i + 1:]
-                if not next_content.startswith("Scope:")
-            )
 
-            if is_parent:
-                node = parent.add(styled_content, style=None)
-                indent_stack.append((indent, node))
-            else:
-                parent.add(styled_content, style=None)
+        def build_tree(label: str, lines: list):
+            node = Tree(label, style="bold magenta" if lines else "cyan")
+            for line in lines:
+                child_node = node.add(Text(line.content, style="yellow"))
+                child_node.children = build_tree("", line.children).children
+            return node
 
+        for section, lines in spec.sections:
+            section_node = build_tree(f"{section}:", lines)
+            tree.children.append(section_node)
         return tree
 
     def _get_spec(self, verbose: bool = False) -> list[tuple[str, str, int]]:

src/base/spec.rs

@@ -1,35 +1,45 @@
 use crate::prelude::*;
 
 use super::schema::{EnrichedValueType, FieldSchema};
+use pyo3::prelude::*;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 use std::ops::Deref;
 
 /// OutputMode enum for displaying spec info in different granularity
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
 pub enum OutputMode {
     Concise,
     Verbose,
 }
 
-impl OutputMode {
-    pub fn from_str(s: &str) -> Self {
-        match s.to_lowercase().as_str() {
-            "concise" => OutputMode::Concise,
-            "verbose" => OutputMode::Verbose,
-            _ => unreachable!(
-                "Invalid format mode: {}. Expected 'concise' or 'verbose'.",
-                s
-            ),
-        }
-    }
-}
-
 /// Formatting spec per output mode
 pub trait SpecFormatter {
     fn format(&self, mode: OutputMode) -> String;
 }
 
+/// A single line in the rendered spec, with optional scope and children
+#[pyclass(get_all, set_all)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RenderedSpecLine {
+    /// The formatted content of the line (e.g., "Import: name=documents, source=LocalFile")
+    pub content: String,
+    /// The scope name, if applicable (e.g., "documents_1" for ForEach scopes)
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub scope: Option<String>,
+    /// Child lines in the hierarchy
+    pub children: Vec<RenderedSpecLine>,
+}
+
+/// A rendered specification, grouped by sections
+#[pyclass(get_all, set_all)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RenderedSpec {
+    /// List of (section_name, lines) pairs
+    pub sections: Vec<(String, Vec<RenderedSpecLine>)>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]
 pub enum SpecString {
@@ -190,7 +200,7 @@ impl ValueMapping {
 }
 
 impl std::fmt::Display for ValueMapping {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> fmt::Result {
         match self {
             ValueMapping::Constant(v) => write!(
                 f,

src/py/mod.rs

@@ -3,7 +3,7 @@ use crate::prelude::*;
 use crate::base::schema::{FieldSchema, ValueType};
 use crate::base::spec::VectorSimilarityMetric;
 use crate::base::spec::{NamedSpec, ReactiveOpSpec};
-use crate::base::spec::{OutputMode, SpecFormatter};
+use crate::base::spec::{OutputMode, RenderedSpec, RenderedSpecLine, SpecFormatter};
 use crate::execution::query;
 use crate::lib_context::{clear_lib_context, get_auth_registry, init_lib_context};
 use crate::ops::interface::{QueryResult, QueryResults};
@@ -198,76 +198,94 @@ impl Flow {
         })
     }
 
-    #[pyo3(signature = (output_mode="concise"))]
-    pub fn get_spec(&self, output_mode: &str) -> PyResult<Vec<(String, String, u32)>> {
-        let mode = OutputMode::from_str(output_mode);
+    #[pyo3(signature = (output_mode=None))]
+    pub fn get_spec(&self, output_mode: Option<Pythonized<OutputMode>>) -> PyResult<RenderedSpec> {
+        let mode = output_mode.map_or(OutputMode::Concise, |m| m.into_inner());
         let spec = &self.0.flow.flow_instance;
-        let mut result = Vec::with_capacity(
-            1 + spec.import_ops.len()
-                + spec.reactive_ops.len()
-                + spec.export_ops.len()
-                + spec.declarations.len(),
+        let mut sections: IndexMap<String, Vec<RenderedSpecLine>> = IndexMap::new();
+
+        // Initialize sections
+        sections.insert(
+            "Header".to_string(),
+            vec![RenderedSpecLine {
+                content: format!("Flow: {}", spec.name),
+                scope: None,
+                children: vec![],
+            }],
         );
-
-        fn extend_with<T, I, F>(
-            out: &mut Vec<(String, String, u32)>,
-            label: &'static str,
-            items: I,
-            f: F,
-        ) where
-            I: IntoIterator<Item = T>,
-            F: Fn(&T) -> String,
-        {
-            out.extend(items.into_iter().map(|item| (label.into(), f(&item), 0)));
+        for key in ["Sources", "Processing", "Targets", "Declarations"] {
+            sections.insert(key.to_string(), Vec::new());
         }
 
-        // Header
-        result.push(("Header".into(), format!("Flow: {}", spec.name), 0));
-
         // Sources
-        extend_with(&mut result, "Sources", spec.import_ops.iter(), |op| {
-            format!("Import: name={}, {}", op.name, op.spec.format(mode))
-        });
+        for op in &spec.import_ops {
+            sections
+                .entry("Sources".to_string())
+                .or_default()
+                .push(RenderedSpecLine {
+                    content: format!("Import: name={}, {}", op.name, op.spec.format(mode)),
+                    scope: None,
+                    children: vec![],
+                });
+        }
 
         // Processing
         fn walk(
             op: &NamedSpec<ReactiveOpSpec>,
-            indent: u32,
             mode: OutputMode,
-            out: &mut Vec<(String, String, u32)>,
-        ) {
-            out.push((
-                "Processing".into(),
-                format!("{}: {}", op.name, op.spec.format(mode)),
-                indent,
-            ));
+            scope: Option<String>,
+        ) -> RenderedSpecLine {
+            let content = format!("{}: {}", op.name, op.spec.format(mode));
+            let mut line = RenderedSpecLine {
+                content,
+                scope,
+                children: vec![],
+            };
 
             if let ReactiveOpSpec::ForEach(fe) = &op.spec {
-                out.push(("Processing".into(), fe.op_scope.to_string(), indent + 1));
                 for nested in &fe.op_scope.ops {
-                    walk(nested, indent + 2, mode, out);
+                    line.children
+                        .push(walk(nested, mode, Some(fe.op_scope.name.clone())));
                 }
             }
+
+            line
         }
 
         for op in &spec.reactive_ops {
-            walk(op, 0, mode, &mut result);
+            sections
+                .entry("Processing".to_string())
+                .or_default()
+                .push(walk(op, mode, None));
         }
 
         // Targets
-        extend_with(&mut result, "Targets", spec.export_ops.iter(), |op| {
-            format!("Export: name={}, {}", op.name, op.spec.format(mode))
-        });
+        for op in &spec.export_ops {
+            sections
+                .entry("Targets".to_string())
+                .or_default()
+                .push(RenderedSpecLine {
+                    content: format!("Export: name={}, {}", op.name, op.spec.format(mode)),
+                    scope: None,
+                    children: vec![],
+                });
+        }
 
         // Declarations
-        extend_with(
-            &mut result,
-            "Declarations",
-            spec.declarations.iter(),
-            |decl| format!("Declaration: {}", decl.format(mode)),
-        );
+        for decl in &spec.declarations {
+            sections
+                .entry("Declarations".to_string())
+                .or_default()
+                .push(RenderedSpecLine {
+                    content: format!("Declaration: {}", decl.format(mode)),
+                    scope: None,
+                    children: vec![],
+                });
+        }
 
-        Ok(result)
+        Ok(RenderedSpec {
+            sections: sections.into_iter().collect(),
+        })
     }
 
     pub fn get_schema(&self) -> Vec<(String, String, String)> {
@@ -518,6 +536,8 @@ fn cocoindex_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<SimpleSemanticsQueryHandler>()?;
     m.add_class::<SetupStatusCheck>()?;
     m.add_class::<PyOpArgSchema>()?;
+    m.add_class::<RenderedSpec>()?;
+    m.add_class::<RenderedSpecLine>()?;
 
     Ok(())
 }