madhurjindal commited on
Commit
837ceb8
1 Parent(s): 7991dd4

Create render_mermaid.py

Browse files
Files changed (1) hide show
  1. 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)