|
1
|
"""Mermaid rendering and chart reproduction utilities.""" |
|
2
|
|
|
3
|
import logging |
|
4
|
from pathlib import Path |
|
5
|
from typing import Dict |
|
6
|
|
|
7
|
logger = logging.getLogger(__name__) |
|
8
|
|
|
9
|
|
|
10
|
def render_mermaid(mermaid_code: str, output_dir: str | Path, name: str) -> Dict[str, Path]: |
|
11
|
""" |
|
12
|
Render mermaid code to SVG and PNG files. |
|
13
|
|
|
14
|
Writes {name}.mermaid (source), {name}.svg, and {name}.png. |
|
15
|
Uses mermaid-py if available, falls back gracefully. |
|
16
|
|
|
17
|
Returns dict with keys: mermaid, svg, png (Paths to generated files). |
|
18
|
""" |
|
19
|
output_dir = Path(output_dir) |
|
20
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
21
|
result: Dict[str, Path] = {} |
|
22
|
|
|
23
|
# Always write source |
|
24
|
mermaid_path = output_dir / f"{name}.mermaid" |
|
25
|
mermaid_path.write_text(mermaid_code) |
|
26
|
result["mermaid"] = mermaid_path |
|
27
|
|
|
28
|
try: |
|
29
|
import mermaid as mmd |
|
30
|
from mermaid.graph import Graph |
|
31
|
|
|
32
|
graph = Graph("diagram", mermaid_code) |
|
33
|
rendered = mmd.Mermaid(graph) |
|
34
|
|
|
35
|
# SVG |
|
36
|
svg_path = output_dir / f"{name}.svg" |
|
37
|
svg_content = rendered.svg_response |
|
38
|
if svg_content: |
|
39
|
if isinstance(svg_content, bytes): |
|
40
|
svg_path.write_bytes(svg_content) |
|
41
|
else: |
|
42
|
svg_path.write_text(svg_content) |
|
43
|
result["svg"] = svg_path |
|
44
|
|
|
45
|
# PNG |
|
46
|
png_path = output_dir / f"{name}.png" |
|
47
|
png_content = rendered.img_response |
|
48
|
if png_content: |
|
49
|
if isinstance(png_content, bytes): |
|
50
|
png_path.write_bytes(png_content) |
|
51
|
else: |
|
52
|
png_path.write_bytes( |
|
53
|
png_content.encode() if isinstance(png_content, str) else png_content |
|
54
|
) |
|
55
|
result["png"] = png_path |
|
56
|
|
|
57
|
except ImportError: |
|
58
|
logger.warning( |
|
59
|
"mermaid-py not installed, skipping SVG/PNG rendering. " |
|
60
|
"Install with: pip install mermaid-py" |
|
61
|
) |
|
62
|
except Exception as e: |
|
63
|
logger.warning(f"Mermaid rendering failed for '{name}': {e}") |
|
64
|
|
|
65
|
return result |
|
66
|
|
|
67
|
|
|
68
|
def reproduce_chart( |
|
69
|
chart_data: dict, |
|
70
|
output_dir: str | Path, |
|
71
|
name: str, |
|
72
|
) -> Dict[str, Path]: |
|
73
|
""" |
|
74
|
Reproduce a chart from extracted data using matplotlib. |
|
75
|
|
|
76
|
chart_data should contain: labels, values, chart_type (bar/line/pie/scatter). |
|
77
|
Returns dict with keys: svg, png (Paths to generated files). |
|
78
|
""" |
|
79
|
output_dir = Path(output_dir) |
|
80
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
81
|
result: Dict[str, Path] = {} |
|
82
|
|
|
83
|
labels = chart_data.get("labels", []) |
|
84
|
values = chart_data.get("values", []) |
|
85
|
chart_type = chart_data.get("chart_type", "bar") |
|
86
|
|
|
87
|
if not labels or not values: |
|
88
|
logger.warning(f"Insufficient chart data for '{name}': missing labels or values") |
|
89
|
return result |
|
90
|
|
|
91
|
try: |
|
92
|
import matplotlib |
|
93
|
|
|
94
|
matplotlib.use("Agg") # Non-interactive backend |
|
95
|
import matplotlib.pyplot as plt |
|
96
|
|
|
97
|
fig, ax = plt.subplots(figsize=(10, 6)) |
|
98
|
|
|
99
|
if chart_type == "bar": |
|
100
|
ax.bar(labels, values) |
|
101
|
elif chart_type == "line": |
|
102
|
ax.plot(labels, values, marker="o") |
|
103
|
elif chart_type == "pie": |
|
104
|
ax.pie(values, labels=labels, autopct="%1.1f%%") |
|
105
|
elif chart_type == "scatter": |
|
106
|
ax.scatter(range(len(values)), values) |
|
107
|
if labels: |
|
108
|
ax.set_xticks(range(len(labels))) |
|
109
|
ax.set_xticklabels(labels, rotation=45, ha="right") |
|
110
|
else: |
|
111
|
ax.bar(labels, values) |
|
112
|
|
|
113
|
if chart_type != "pie": |
|
114
|
ax.set_xlabel("") |
|
115
|
ax.set_ylabel("") |
|
116
|
plt.xticks(rotation=45, ha="right") |
|
117
|
|
|
118
|
plt.tight_layout() |
|
119
|
|
|
120
|
# SVG |
|
121
|
svg_path = output_dir / f"{name}_chart.svg" |
|
122
|
fig.savefig(svg_path, format="svg") |
|
123
|
result["svg"] = svg_path |
|
124
|
|
|
125
|
# PNG |
|
126
|
png_path = output_dir / f"{name}_chart.png" |
|
127
|
fig.savefig(png_path, format="png", dpi=150) |
|
128
|
result["png"] = png_path |
|
129
|
|
|
130
|
plt.close(fig) |
|
131
|
|
|
132
|
except ImportError: |
|
133
|
logger.warning("matplotlib not installed, skipping chart reproduction") |
|
134
|
except Exception as e: |
|
135
|
logger.warning(f"Chart reproduction failed for '{name}': {e}") |
|
136
|
|
|
137
|
return result |
|
138
|
|