File size: 5,587 Bytes
3943768 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import argparse
import os
import subprocess
import tempfile
import datetime
import random
import string
import shlex
import uuid
def generate_unique_filename(format):
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
random_string = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
return f"mermaid_{timestamp}_{random_string}.{format}"
def find_chrome_path():
home_dir = os.path.expanduser("~")
cache_dir = os.path.join(home_dir, ".cache", "puppeteer")
try:
cmd = f"find {shlex.quote(cache_dir)} -name chrome-headless-shell -type f | sort -V | tail -n 1"
result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
chrome_path = result.stdout.strip()
if not chrome_path:
print("Chrome headless shell not found in the expected location.")
return None
return chrome_path
except subprocess.CalledProcessError as e:
print(f"An error occurred while trying to find Chrome: {e}")
return None
def render_mermaid(mermaid_code, output_file, format='svg'):
# Find Chrome path
use_headless = False
if use_headless:
chrome_path = find_chrome_path()
if not chrome_path:
raise Exception("Chrome headless shell not found. Unable to render Mermaid diagram.")
# Set PUPPETEER_EXECUTABLE_PATH environment variable
os.environ["PUPPETEER_EXECUTABLE_PATH"] = chrome_path
# else let it default to chromium-browser, just still requires no sandbox
elif os.path.isfile('/usr/bin/chromium-browser'):
os.environ["PUPPETEER_EXECUTABLE_PATH"] = '/usr/bin/chromium-browser'
# Create a temporary file for the Mermaid code
with tempfile.NamedTemporaryFile(mode='w', suffix='.mmd', delete=False) as temp:
temp.write(mermaid_code)
temp_path = temp.name
config_file = f'puppeteer-config{str(uuid.uuid4())}.json'
try:
# Construct the mmdc command
with open(config_file, 'wt') as f:
f.write('{"args": ["--no-sandbox"]}')
cmd = ['mmdc', '-i', temp_path, '-o', output_file, '-f', format, '-p', config_file]
# Run the mmdc command
result = subprocess.run(cmd, check=True, capture_output=True, text=True)
# Check if there was any output (warnings, etc.)
if result.stdout:
print("mmdc output:", result.stdout)
if result.stderr:
print("mmdc warnings/errors:", result.stderr)
print(f"Created output file in {format} format: {output_file}")
# Always make PNG version too, hard for other tools to svg -> png
if format != 'png':
# Construct the mmdc command
base_name = '.'.join(output_file.split('.')[:-1])
output_file_png = base_name + '.png'
# FIXME: Would be best to optimize for aspect ratio in choosing -w or -H
cmd = ['mmdc', '-i', temp_path, '-o', output_file_png, '-f', 'png', '-w', '2048', '-p', config_file]
# Run the mmdc command
result = subprocess.run(cmd, check=True, capture_output=True, text=True)
# Check if there was any output (warnings, etc.)
if result.stdout:
print("mmdc output:", result.stdout)
if result.stderr:
print("mmdc warnings/errors:", result.stderr)
print(
f"Created mermaid output file in PNG format: {output_file_png} that is a conversion of {output_file}. "
"Use this for image_query to analyze what SVG looks like, "
"because other tools do not retain fonts when making PNG."
)
# Return the full path of the output file
return os.path.abspath(output_file)
finally:
# Clean up the temporary file
os.unlink(temp_path)
if os.path.isfile(config_file):
try:
os.remove(config_file)
except FileNotFoundError:
pass
def main():
parser = argparse.ArgumentParser(description='Render Mermaid diagrams from a file or direct input using mmdc.')
input_group = parser.add_mutually_exclusive_group(required=True)
input_group.add_argument('-f', '--file', '--input', help='Input file containing Mermaid code')
input_group.add_argument('-c', '--code', help='Direct Mermaid code input', nargs='+')
parser.add_argument('-o', '--output', help='Output file name (default: auto-generated unique name)')
args = parser.parse_args()
# If no output file is specified, create a unique name
if args.output is None:
format = 'svg'
args.output = generate_unique_filename(format)
else:
format = args.output.split('.')[-1]
assert format in ['svg', 'png', 'pdf'], f"Invalid output filename {args.output} with format: {format}"
try:
# Determine the Mermaid code source
if args.file:
with open(args.file, 'r') as f:
mermaid_code = f.read()
else:
mermaid_code = ' '.join(args.code)
# Render the diagram and get the full path of the output file
output_path = render_mermaid(mermaid_code, args.output, format=format)
print(f"Mermaid diagram rendered successfully.")
print(f"Output file: {output_path}")
except subprocess.CalledProcessError as e:
print(f"Error rendering Mermaid diagram: {e}")
print(f"mmdc output: {e.output}")
print(f"mmdc error: {e.stderr}")
if __name__ == "__main__":
main()
|