MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/9/4 16:12
5
@Author  : alitrack
6
@File    : mmdc_pyppeteer.py
7
"""
8
import os
9
from urllib.parse import urljoin
10

11
from pyppeteer import launch
12

13
from metagpt.config2 import config
14
from metagpt.logs import logger
15

16

17
async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int:
18
    """
19
    Converts the given Mermaid code to various output formats and saves them to files.
20

21
    Args:
22
        mermaid_code (str): The Mermaid code to convert.
23
        output_file_without_suffix (str): The output file name without the file extension.
24
        width (int, optional): The width of the output image in pixels. Defaults to 2048.
25
        height (int, optional): The height of the output image in pixels. Defaults to 2048.
26

27
    Returns:
28
        int: Returns 1 if the conversion and saving were successful, -1 otherwise.
29
    """
30
    suffixes = ["png", "svg", "pdf"]
31
    __dirname = os.path.dirname(os.path.abspath(__file__))
32

33
    if config.mermaid.pyppeteer_path:
34
        browser = await launch(
35
            headless=True,
36
            executablePath=config.mermaid.pyppeteer_path,
37
            args=["--disable-extensions", "--no-sandbox"],
38
        )
39
    else:
40
        logger.error("Please set the var mermaid.pyppeteer_path in the config2.yaml.")
41
        return -1
42
    page = await browser.newPage()
43
    device_scale_factor = 1.0
44

45
    async def console_message(msg):
46
        logger.info(msg.text)
47

48
    page.on("console", console_message)
49

50
    try:
51
        await page.setViewport(viewport={"width": width, "height": height, "deviceScaleFactor": device_scale_factor})
52

53
        mermaid_html_path = os.path.abspath(os.path.join(__dirname, "index.html"))
54
        mermaid_html_url = urljoin("file:", mermaid_html_path)
55
        await page.goto(mermaid_html_url)
56

57
        await page.querySelector("div#container")
58
        # mermaid_config = {}
59
        background_color = "#ffffff"
60
        # my_css = ""
61
        await page.evaluate(f'document.body.style.background = "{background_color}";')
62

63
        # metadata = await page.evaluate(
64
        #     """async ([definition, mermaidConfig, myCSS, backgroundColor]) => {
65
        #     const { mermaid, zenuml } = globalThis;
66
        #     await mermaid.registerExternalDiagrams([zenuml]);
67
        #     mermaid.initialize({ startOnLoad: false, ...mermaidConfig });
68
        #     const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container'));
69
        #     document.getElementById('container').innerHTML = svg;
70
        #     const svgElement = document.querySelector('svg');
71
        #     svgElement.style.backgroundColor = backgroundColor;
72
        #
73
        #     if (myCSS) {
74
        #         const style = document.createElementNS('http://www.w3.org/2000/svg', 'style');
75
        #         style.appendChild(document.createTextNode(myCSS));
76
        #         svgElement.appendChild(style);
77
        #     }
78
        # }""",
79
        #     [mermaid_code, mermaid_config, my_css, background_color],
80
        # )
81

82
        if "svg" in suffixes:
83
            svg_xml = await page.evaluate(
84
                """() => {
85
                const svg = document.querySelector('svg');
86
                const xmlSerializer = new XMLSerializer();
87
                return xmlSerializer.serializeToString(svg);
88
            }"""
89
            )
90
            logger.info(f"Generating {output_file_without_suffix}.svg..")
91
            with open(f"{output_file_without_suffix}.svg", "wb") as f:
92
                f.write(svg_xml.encode("utf-8"))
93

94
        if "png" in suffixes:
95
            clip = await page.evaluate(
96
                """() => {
97
                const svg = document.querySelector('svg');
98
                const rect = svg.getBoundingClientRect();
99
                return {
100
                    x: Math.floor(rect.left),
101
                    y: Math.floor(rect.top),
102
                    width: Math.ceil(rect.width),
103
                    height: Math.ceil(rect.height)
104
                };
105
            }"""
106
            )
107
            await page.setViewport(
108
                {
109
                    "width": clip["x"] + clip["width"],
110
                    "height": clip["y"] + clip["height"],
111
                    "deviceScaleFactor": device_scale_factor,
112
                }
113
            )
114
            screenshot = await page.screenshot(clip=clip, omit_background=True, scale="device")
115
            logger.info(f"Generating {output_file_without_suffix}.png..")
116
            with open(f"{output_file_without_suffix}.png", "wb") as f:
117
                f.write(screenshot)
118
        if "pdf" in suffixes:
119
            pdf_data = await page.pdf(scale=device_scale_factor)
120
            logger.info(f"Generating {output_file_without_suffix}.pdf..")
121
            with open(f"{output_file_without_suffix}.pdf", "wb") as f:
122
                f.write(pdf_data)
123
        return 0
124
    except Exception as e:
125
        logger.error(e)
126
        return -1
127
    finally:
128
        await browser.close()
129