refactor: update spec formatting logic to use SpecFormatter trait

Author: lemorage

python/cocoindex/flow.py

@@ -493,7 +493,7 @@ def _render_spec(self, verbose: bool = False) -> Tree:
         return tree
 
     def _get_spec(self, verbose: bool = False) -> list[tuple[str, str, int]]:
-        return self._lazy_engine_flow().get_spec(format_mode="verbose" if verbose else "concise")
+        return self._lazy_engine_flow().get_spec(output_mode="verbose" if verbose else "concise")
 
     def _get_schema(self) -> list[tuple[str, str, str]]:
         return self._lazy_engine_flow().get_schema()

src/base/spec.rs

@@ -2,27 +2,34 @@ use crate::prelude::*;
 
 use super::schema::{EnrichedValueType, FieldSchema};
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
 use std::fmt;
 use std::ops::Deref;
 
-// Define SpecFormatMode enum for type-safe formatting
+/// OutputMode enum for displaying spec info in different granularity
 #[derive(Debug, Clone, Copy, PartialEq)]
-pub enum SpecFormatMode {
+pub enum OutputMode {
     Concise,
     Verbose,
 }
 
-impl SpecFormatMode {
-    pub fn from_str(s: &str) -> Result<Self, String> {
+impl OutputMode {
+    pub fn from_str(s: &str) -> Self {
         match s.to_lowercase().as_str() {
-            "concise" => Ok(SpecFormatMode::Concise),
-            "verbose" => Ok(SpecFormatMode::Verbose),
-            _ => Err(format!("Invalid format mode: {}", s)),
+            "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;
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]
 pub enum SpecString {
@@ -233,69 +240,31 @@ pub struct OpSpec {
     pub spec: serde_json::Map<String, serde_json::Value>,
 }
 
-impl OpSpec {
-    pub fn format_concise(&self) -> String {
-        let mut parts = vec![];
-        for (key, value) in self.spec.iter() {
-            match value {
-                Value::String(s) => parts.push(format!("{}={}", key, s)),
-                Value::Array(arr) => {
-                    let items = arr
-                        .iter()
-                        .filter_map(|v| v.as_str())
-                        .collect::<Vec<_>>()
-                        .join(",");
-                    if !items.is_empty() {
-                        parts.push(format!("{}={}", key, items));
-                    }
-                }
-                Value::Object(obj) => {
-                    if let Some(model) = obj.get("model").and_then(|v| v.as_str()) {
-                        parts.push(format!("{}={}", key, model));
-                    }
-                }
-                _ => {}
-            }
-        }
-        if parts.is_empty() {
-            self.kind.clone()
-        } else {
-            format!("{}({})", self.kind, parts.join(", "))
-        }
-    }
-
-    pub fn format_verbose(&self) -> String {
-        let spec_str = serde_json::to_string_pretty(&self.spec)
-            .map(|s| {
-                let lines: Vec<&str> = s.lines().collect();
-                if lines.len() < s.lines().count() {
-                    lines
-                        .into_iter()
-                        .chain(["..."])
-                        .collect::<Vec<_>>()
-                        .join("\n  ")
-                } else {
-                    lines.join("\n  ")
-                }
-            })
-            .unwrap_or("#serde_error".to_string());
-        format!("{}({})", self.kind, spec_str)
-    }
-
-    pub fn format(&self, mode: SpecFormatMode) -> String {
+impl SpecFormatter for OpSpec {
+    fn format(&self, mode: OutputMode) -> String {
         match mode {
-            SpecFormatMode::Concise => self.format_concise(),
-            SpecFormatMode::Verbose => self.format_verbose(),
+            OutputMode::Concise => self.kind.clone(),
+            OutputMode::Verbose => {
+                let spec_str = serde_json::to_string_pretty(&self.spec)
+                    .map(|s| {
+                        let lines: Vec<&str> = s.lines().collect();
+                        if lines.len() < s.lines().count() {
+                            lines
+                                .into_iter()
+                                .chain(["..."])
+                                .collect::<Vec<_>>()
+                                .join("\n  ")
+                        } else {
+                            lines.join("\n  ")
+                        }
+                    })
+                    .unwrap_or("#serde_error".to_string());
+                format!("{}({})", self.kind, spec_str)
+            }
         }
     }
 }
 
-impl fmt::Display for OpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format_concise())
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct SourceRefreshOptions {
     pub refresh_interval: Option<std::time::Duration>,
@@ -319,19 +288,16 @@ pub struct ImportOpSpec {
     pub refresh_options: SourceRefreshOptions,
 }
 
-impl ImportOpSpec {
-    pub fn format(&self, mode: SpecFormatMode) -> String {
-        let source = match mode {
-            SpecFormatMode::Concise => self.source.format_concise(),
-            SpecFormatMode::Verbose => self.source.format_verbose(),
-        };
+impl SpecFormatter for ImportOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
+        let source = self.source.format(mode);
         format!("source={}, refresh={}", source, self.refresh_options)
     }
 }
 
 impl fmt::Display for ImportOpSpec {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format(SpecFormatMode::Concise))
+        write!(f, "{}", self.format(OutputMode::Concise))
     }
 }
 
@@ -341,31 +307,22 @@ pub struct TransformOpSpec {
     pub op: OpSpec,
 }
 
-impl TransformOpSpec {
-    pub fn format(&self, mode: SpecFormatMode) -> String {
+impl SpecFormatter for TransformOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
         let inputs = self
             .inputs
             .iter()
             .map(ToString::to_string)
             .collect::<Vec<_>>()
             .join(",");
-        let op_str = match mode {
-            SpecFormatMode::Concise => self.op.format_concise(),
-            SpecFormatMode::Verbose => self.op.format_verbose(),
-        };
+        let op_str = self.op.format(mode);
         match mode {
-            SpecFormatMode::Concise => format!("op={}, inputs={}", op_str, inputs),
-            SpecFormatMode::Verbose => format!("op={}, inputs=[{}]", op_str, inputs),
+            OutputMode::Concise => format!("op={}, inputs={}", op_str, inputs),
+            OutputMode::Verbose => format!("op={}, inputs=[{}]", op_str, inputs),
         }
     }
 }
 
-impl fmt::Display for TransformOpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format(SpecFormatMode::Concise))
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ForEachOpSpec {
     /// Mapping that provides a table to apply reactive operations to.
@@ -377,21 +334,17 @@ impl ForEachOpSpec {
     pub fn get_label(&self) -> String {
         format!("Loop over {}", self.field_path)
     }
+}
 
-    pub fn format(&self, mode: SpecFormatMode) -> String {
+impl SpecFormatter for ForEachOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
         match mode {
-            SpecFormatMode::Concise => self.get_label(),
-            SpecFormatMode::Verbose => format!("field={}", self.field_path),
+            OutputMode::Concise => self.get_label(),
+            OutputMode::Verbose => format!("field={}", self.field_path),
         }
     }
 }
 
-impl fmt::Display for ForEachOpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "field={}", self.field_path)
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CollectOpSpec {
     /// Field values to be collected.
@@ -405,17 +358,17 @@ pub struct CollectOpSpec {
     pub auto_uuid_field: Option<FieldName>,
 }
 
-impl CollectOpSpec {
-    pub fn format(&self, mode: SpecFormatMode) -> String {
+impl SpecFormatter for CollectOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
         let uuid = self.auto_uuid_field.as_deref().unwrap_or("none");
         match mode {
-            SpecFormatMode::Concise => {
+            OutputMode::Concise => {
                 format!(
                     "collector={}, input={}, uuid={}",
                     self.collector_name, self.input, uuid
                 )
             }
-            SpecFormatMode::Verbose => {
+            OutputMode::Verbose => {
                 format!(
                     "scope={}, collector={}, input=[{}], uuid={}",
                     self.scope_name, self.collector_name, self.input, uuid
@@ -425,12 +378,6 @@ impl CollectOpSpec {
     }
 }
 
-impl fmt::Display for CollectOpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format(SpecFormatMode::Concise))
-    }
-}
-
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
 pub enum VectorSimilarityMetric {
     CosineSimilarity,
@@ -493,29 +440,20 @@ pub struct ExportOpSpec {
     pub setup_by_user: bool,
 }
 
-impl ExportOpSpec {
-    pub fn format(&self, mode: SpecFormatMode) -> String {
-        let target_str = match mode {
-            SpecFormatMode::Concise => self.target.format_concise(),
-            SpecFormatMode::Verbose => self.target.format_verbose(),
-        };
+impl SpecFormatter for ExportOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
+        let target_str = self.target.format(mode);
         let base = format!(
             "collector={}, target={}, {}",
             self.collector_name, target_str, self.index_options
         );
         match mode {
-            SpecFormatMode::Concise => base,
-            SpecFormatMode::Verbose => format!("{}, setup_by_user={}", base, self.setup_by_user),
+            OutputMode::Concise => base,
+            OutputMode::Verbose => format!("{}, setup_by_user={}", base, self.setup_by_user),
         }
     }
 }
 
-impl fmt::Display for ExportOpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format(SpecFormatMode::Concise))
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "action")]
 pub enum ReactiveOpSpec {
@@ -524,25 +462,19 @@ pub enum ReactiveOpSpec {
     Collect(CollectOpSpec),
 }
 
-impl ReactiveOpSpec {
-    pub fn format(&self, mode: SpecFormatMode) -> String {
+impl SpecFormatter for ReactiveOpSpec {
+    fn format(&self, mode: OutputMode) -> String {
         match self {
             ReactiveOpSpec::Transform(t) => format!("Transform: {}", t.format(mode)),
             ReactiveOpSpec::ForEach(fe) => match mode {
-                SpecFormatMode::Concise => format!("{}", fe.get_label()),
-                SpecFormatMode::Verbose => format!("ForEach: {}", fe.format(mode)),
+                OutputMode::Concise => format!("{}", fe.get_label()),
+                OutputMode::Verbose => format!("ForEach: {}", fe.format(mode)),
             },
             ReactiveOpSpec::Collect(c) => format!("Collect: {}", c.format(mode)),
         }
     }
 }
 
-impl fmt::Display for ReactiveOpSpec {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.format(SpecFormatMode::Concise))
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ReactiveOpScope {
     pub name: ScopeName,

src/py/mod.rs

@@ -1,9 +1,9 @@
 use crate::prelude::*;
 
 use crate::base::schema::{FieldSchema, ValueType};
-use crate::base::spec::SpecFormatMode;
 use crate::base::spec::VectorSimilarityMetric;
 use crate::base::spec::{NamedSpec, ReactiveOpSpec};
+use crate::base::spec::{OutputMode, 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,10 +198,9 @@ impl Flow {
         })
     }
 
-    #[pyo3(signature = (format_mode="concise"))]
-    pub fn get_spec(&self, format_mode: &str) -> PyResult<Vec<(String, String, u32)>> {
-        let mode = SpecFormatMode::from_str(format_mode)
-            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
+    #[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);
         let spec = &self.0.flow.flow_instance;
         let mut result = Vec::with_capacity(
             1 + spec.import_ops.len()
@@ -234,7 +233,7 @@ impl Flow {
         fn walk(
             op: &NamedSpec<ReactiveOpSpec>,
             indent: u32,
-            mode: SpecFormatMode,
+            mode: OutputMode,
             out: &mut Vec<(String, String, u32)>,
         ) {
             out.push((