Skip to content

feat(pyo3): implement flow spec pretty print and add verbose mode #442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 10, 2025
13 changes: 7 additions & 6 deletions python/cocoindex/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,16 +55,17 @@ def ls(show_all: bool):

@cli.command()
@click.argument("flow_name", type=str, required=False)
@click.option("--color/--no-color", default=True)
def show(flow_name: str | None, color: bool):
@click.option("--color/--no-color", default=True, help="Enable or disable colored output.")
@click.option("--verbose", is_flag=True, help="Show verbose output with full details.")
def show(flow_name: str | None, color: bool, verbose: bool):
"""
Show the flow spec in a readable format with colored output,
including the schema.
Show the flow spec and schema in a readable format with colored output.
"""
flow = _flow_by_name(flow_name)
console = Console(no_color=not color)
console.print(flow._render_text())
console.print(flow._render_spec(verbose=verbose))

console.print()
table = Table(
title=f"Schema for Flow: {flow.name}",
show_header=True,
Expand All @@ -74,7 +75,7 @@ def show(flow_name: str | None, color: bool):
table.add_column("Type", style="green")
table.add_column("Attributes", style="yellow")

for field_name, field_type, attr_str in flow._render_schema():
for field_name, field_type, attr_str in flow._get_schema():
table.add_row(field_name, field_type, attr_str)

console.print(table)
Expand Down
89 changes: 40 additions & 49 deletions python/cocoindex/flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from enum import Enum
from dataclasses import dataclass
from rich.text import Text
from rich.console import Console
from rich.tree import Tree

from . import _engine
from . import index
Expand Down Expand Up @@ -454,61 +454,52 @@ def _lazy_engine_flow() -> _engine.Flow:
return engine_flow
self._lazy_engine_flow = _lazy_engine_flow

def _format_flow(self, flow_dict: dict) -> Text:
output = Text()
def _render_spec(self, verbose: bool = False) -> Tree:
"""
Render the flow spec as a styled rich Tree with hierarchical structure.
"""
tree = Tree(f"Flow: {self.name}", style="cyan")
current_section = None
section_node = None
indent_stack = []

def add_line(content, indent=0, style=None, end="\n"):
output.append(" " * indent)
output.append(content, style=style)
output.append(end)
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

def format_key_value(key, value, indent):
if isinstance(value, (dict, list)):
add_line(f"- {key}:", indent, style="green")
format_data(value, indent + 2)
else:
add_line(f"- {key}:", indent, style="green", end="")
add_line(f" {value}", style="yellow")

def format_data(data, indent=0):
if isinstance(data, dict):
for key, value in data.items():
format_key_value(key, value, indent)
elif isinstance(data, list):
for i, item in enumerate(data):
format_key_value(f"[{i}]", item, indent)
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:
add_line(str(data), indent, style="yellow")

# Header
flow_name = flow_dict.get("name", "Unnamed")
add_line(f"Flow: {flow_name}", style="bold cyan")

# Section
for section_title, section_key in [
("Sources:", "import_ops"),
("Processing:", "reactive_ops"),
("Targets:", "export_ops"),
]:
add_line("")
add_line(section_title, style="bold cyan")
format_data(flow_dict.get(section_key, []), indent=0)

return output

def _render_text(self) -> Text:
flow_spec_str = str(self._lazy_engine_flow())
try:
flow_dict = json.loads(flow_spec_str)
return self._format_flow(flow_dict)
except json.JSONDecodeError:
return Text(flow_spec_str)
parent.add(styled_content, style=None)

return tree

def _get_spec(self, verbose: bool = False) -> list[tuple[str, str, int]]:
return self._lazy_engine_flow().get_spec(output_mode="verbose" if verbose else "concise")

def _render_schema(self) -> list[tuple[str, str, str]]:
def _get_schema(self) -> list[tuple[str, str, str]]:
return self._lazy_engine_flow().get_schema()

def __str__(self):
return str(self._render_text())
return str(self._get_spec())

def __repr__(self):
return repr(self._lazy_engine_flow())
Expand Down
Loading