madhurjindal
commited on
Commit
•
837ceb8
1
Parent(s):
7991dd4
Create render_mermaid.py
Browse files- render_mermaid.py +102 -0
render_mermaid.py
ADDED
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
import os
|
3 |
+
import sys
|
4 |
+
import time
|
5 |
+
from selenium import webdriver
|
6 |
+
from selenium.webdriver.chrome.options import Options
|
7 |
+
from selenium.webdriver.chrome.service import Service # Import Service
|
8 |
+
from webdriver_manager.chrome import ChromeDriverManager
|
9 |
+
|
10 |
+
|
11 |
+
def render_mermaid(mermaid_code, output_path):
|
12 |
+
# Create an HTML file with the Mermaid code
|
13 |
+
mermaid_html = f"""
|
14 |
+
<!DOCTYPE html>
|
15 |
+
<html>
|
16 |
+
<head>
|
17 |
+
<meta charset="utf-8">
|
18 |
+
<script src="https://unpkg.com/mermaid@9/dist/mermaid.min.js"></script>
|
19 |
+
</head>
|
20 |
+
<body>
|
21 |
+
<div class="mermaid">
|
22 |
+
{mermaid_code}
|
23 |
+
</div>
|
24 |
+
<script>
|
25 |
+
mermaid.initialize({{ startOnLoad: true }});
|
26 |
+
</script>
|
27 |
+
</body>
|
28 |
+
</html>
|
29 |
+
"""
|
30 |
+
|
31 |
+
# Write the HTML content to a temporary file
|
32 |
+
temp_html_file = "temp_mermaid.html"
|
33 |
+
with open(temp_html_file, "w", encoding="utf-8") as file:
|
34 |
+
file.write(mermaid_html)
|
35 |
+
|
36 |
+
# Set up Selenium with headless Chrome
|
37 |
+
chrome_options = Options()
|
38 |
+
chrome_options.add_argument("--headless")
|
39 |
+
chrome_options.add_argument("--disable-gpu") # Add this line for compatibility
|
40 |
+
chrome_options.add_argument("--no-sandbox") # Add this line if running as root
|
41 |
+
chrome_options.add_argument(
|
42 |
+
"--disable-dev-shm-usage"
|
43 |
+
) # Overcome limited resource problems
|
44 |
+
|
45 |
+
# Initialize the WebDriver using Service
|
46 |
+
service = Service(ChromeDriverManager().install())
|
47 |
+
driver = webdriver.Chrome(service=service, options=chrome_options)
|
48 |
+
|
49 |
+
try:
|
50 |
+
# Open the HTML file in the browser
|
51 |
+
driver.get("file://" + os.path.abspath(temp_html_file))
|
52 |
+
|
53 |
+
# Wait for the diagram to render
|
54 |
+
time.sleep(2) # Increase if necessary
|
55 |
+
|
56 |
+
# Find the diagram element
|
57 |
+
element = driver.find_element("css selector", ".mermaid")
|
58 |
+
|
59 |
+
# Take a screenshot of the element
|
60 |
+
element.screenshot(output_path)
|
61 |
+
|
62 |
+
except Exception as e:
|
63 |
+
print(f"Error during rendering: {e}")
|
64 |
+
raise e
|
65 |
+
|
66 |
+
finally:
|
67 |
+
# Close the browser and clean up
|
68 |
+
driver.quit()
|
69 |
+
os.remove(temp_html_file)
|
70 |
+
|
71 |
+
|
72 |
+
if __name__ == "__main__":
|
73 |
+
parser = argparse.ArgumentParser(description="Render Mermaid diagram to an image.")
|
74 |
+
parser.add_argument("--code", type=str, help="Mermaid code as a string.")
|
75 |
+
parser.add_argument(
|
76 |
+
"--input", type=str, help="Path to a file containing Mermaid code."
|
77 |
+
)
|
78 |
+
parser.add_argument(
|
79 |
+
"--output", type=str, default="diagram.png", help="Output image file path."
|
80 |
+
)
|
81 |
+
args = parser.parse_args()
|
82 |
+
|
83 |
+
if args.code:
|
84 |
+
mermaid_code = args.code
|
85 |
+
elif args.input:
|
86 |
+
with open(args.input, "r", encoding="utf-8") as f:
|
87 |
+
mermaid_code = f.read()
|
88 |
+
else:
|
89 |
+
# Read from standard input
|
90 |
+
print("Please enter your Mermaid code (Press Ctrl+D to end input):")
|
91 |
+
mermaid_code = sys.stdin.read()
|
92 |
+
if not mermaid_code.strip():
|
93 |
+
print("Error: No Mermaid code provided.")
|
94 |
+
exit(1)
|
95 |
+
|
96 |
+
output_path = args.output
|
97 |
+
try:
|
98 |
+
render_mermaid(mermaid_code, output_path)
|
99 |
+
print(f"Diagram saved to {output_path}")
|
100 |
+
except Exception as e:
|
101 |
+
print(f"Error rendering diagram: {e}")
|
102 |
+
exit(1)
|