From 69ea116d1a265b4b9c0e112832392d5decd7aab1 Mon Sep 17 00:00:00 2001 From: "hy.li" Date: Wed, 13 Sep 2023 12:41:38 +0800 Subject: [PATCH] change to async --- metagpt/actions/design_api.py | 18 +-- metagpt/utils/mermaid.py | 71 ++++++------ metagpt/utils/mmdc_ink.py | 52 ++++----- metagpt/utils/mmdc_playwright.py | 179 +++++++++++++----------------- metagpt/utils/mmdc_pyppeteer.py | 182 ++++++++++++++----------------- 5 files changed, 225 insertions(+), 277 deletions(-) diff --git a/metagpt/actions/design_api.py b/metagpt/actions/design_api.py index abd1f9d4c..4d17e4f5e 100644 --- a/metagpt/actions/design_api.py +++ b/metagpt/actions/design_api.py @@ -103,23 +103,23 @@ class WriteDesign(Action): pass # Folder does not exist, but we don't care workspace.mkdir(parents=True, exist_ok=True) - def _save_prd(self, docs_path, resources_path, prd): + async def _save_prd(self, docs_path, resources_path, prd): prd_file = docs_path / 'prd.md' quadrant_chart = CodeParser.parse_code(block="Competitive Quadrant Chart", text=prd) - mermaid_to_file(quadrant_chart, resources_path / 'competitive_analysis') + await mermaid_to_file(quadrant_chart, resources_path / 'competitive_analysis') logger.info(f"Saving PRD to {prd_file}") prd_file.write_text(prd) - def _save_system_design(self, docs_path, resources_path, content): + async def _save_system_design(self, docs_path, resources_path, content): data_api_design = CodeParser.parse_code(block="Data structures and interface definitions", text=content) seq_flow = CodeParser.parse_code(block="Program call flow", text=content) - mermaid_to_file(data_api_design, resources_path / 'data_api_design') - mermaid_to_file(seq_flow, resources_path / 'seq_flow') + await mermaid_to_file(data_api_design, resources_path / 'data_api_design') + await mermaid_to_file(seq_flow, resources_path / 'seq_flow') system_design_file = docs_path / 'system_design.md' logger.info(f"Saving System Designs to {system_design_file}") system_design_file.write_text(content) - def _save(self, context, system_design): + async def _save(self, context, system_design): if isinstance(system_design, ActionOutput): content = system_design.content ws_name = CodeParser.parse_str(block="Python package name", text=content) @@ -132,13 +132,13 @@ class WriteDesign(Action): resources_path = workspace / 'resources' docs_path.mkdir(parents=True, exist_ok=True) resources_path.mkdir(parents=True, exist_ok=True) - self._save_prd(docs_path, resources_path, context[-1].content) - self._save_system_design(docs_path, resources_path, content) + await self._save_prd(docs_path, resources_path, context[-1].content) + await self._save_system_design(docs_path, resources_path, content) async def run(self, context): prompt = PROMPT_TEMPLATE.format(context=context, format_example=FORMAT_EXAMPLE) # system_design = await self._aask(prompt) system_design = await self._aask_v1(prompt, "system_design", OUTPUT_MAPPING) - self._save(context, system_design) + await self._save(context, system_design) return system_design \ No newline at end of file diff --git a/metagpt/utils/mermaid.py b/metagpt/utils/mermaid.py index b13199a93..d2cce3965 100644 --- a/metagpt/utils/mermaid.py +++ b/metagpt/utils/mermaid.py @@ -2,9 +2,10 @@ # -*- coding: utf-8 -*- """ @Time : 2023/7/4 10:53 -@Author : alexanderwu +@Author : alexanderwu alitrack @File : mermaid.py """ +import asyncio import subprocess from pathlib import Path @@ -15,18 +16,22 @@ from metagpt.utils.common import check_cmd_exists import os import sys -def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int: +async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int: """suffix: png/svg/pdf :param mermaid_code: mermaid code :param output_file_without_suffix: output filename :param width: :param height: - :return: 0 if succed, -1 if failed + :return: 0 if succeed, -1 if failed """ # Write the Mermaid code to a temporary file + dir_name = os.path.dirname(output_file_without_suffix) + if dir_name and not os.path.exists(dir_name): + os.makedirs(dir_name) tmp = Path(f"{output_file_without_suffix}.mmd") tmp.write_text(mermaid_code, encoding="utf-8") + engine = CONFIG.mermaid_engine.lower() if engine == "nodejs": if check_cmd_exists("mmdc") != 0: @@ -39,8 +44,7 @@ def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height logger.info(f"Generating {output_file}..") if CONFIG.puppeteer_config: - subprocess.run( - [ + commands =[ CONFIG.mmdc, "-p", CONFIG.puppeteer_config, @@ -53,33 +57,32 @@ def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height "-H", str(height), ] - ) else: - subprocess.run([CONFIG.mmdc, "-i", str(tmp), "-o", output_file, "-w", str(width), "-H", str(height)]) - else: - if engine not in ['playwright', 'pyppeteer', 'ink']: - logger.warning(f"Unsupported mermaid engine: {engine}") - return -1 - __dirname = os.path.dirname(os.path.abspath(__file__)) - module_path = os.path.join(__dirname, f'mmdc_{engine}.py') - - # 构建命令行参数 - command = [ - sys.executable, - module_path, - "-i",mermaid_code, - "-o",output_file_without_suffix - ] + commands =[CONFIG.mmdc, "-i", str(tmp), "-o", output_file, "-w", str(width), "-H", str(height)] + process = await asyncio.create_subprocess_exec( + *commands, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) - # 执行命令 - try: - result = subprocess.run(command, text=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - logger.info(result.stdout) - if result.stderr: - logger.error(result.stderr) - except subprocess.CalledProcessError as e: - logger.error(f"Command execution failed with return code {e.returncode}") - logger.error(e.output) + stdout, stderr = await process.communicate() + if stdout: + logger.info(stdout.decode()) + if stderr: + logger.error(stderr.decode()) + else: + + if engine =='playwright': + from metagpt.utils.mmdc_playwright import mermaid_to_file + return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height) + elif engine =='pyppeteer': + from metagpt.utils.mmdc_pyppeteer import mermaid_to_file + return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height) + elif engine =='ink': + from metagpt.utils.mmdc_ink import mermaid_to_file + return await mermaid_to_file(mermaid_code, output_file_without_suffix) + else: + logger.warning(f"Unsupported mermaid engine: {engine}") return 0 @@ -134,7 +137,9 @@ MMC2 = """sequenceDiagram SE-->>M: return summary""" + if __name__ == "__main__": - # logger.info(print_members(print_members)) - mermaid_to_file(MMC1, PROJECT_ROOT / "tmp/1.png") - mermaid_to_file(MMC2, PROJECT_ROOT / "tmp/2.png") + loop = asyncio.new_event_loop() + result = loop.run_until_complete(mermaid_to_file(MMC1, PROJECT_ROOT / f"{CONFIG.mermaid_engine}/1")) + result = loop.run_until_complete(mermaid_to_file(MMC2, PROJECT_ROOT / f"{CONFIG.mermaid_engine}/1")) + loop.close() diff --git a/metagpt/utils/mmdc_ink.py b/metagpt/utils/mmdc_ink.py index ce50b11cd..3d91cde9d 100644 --- a/metagpt/utils/mmdc_ink.py +++ b/metagpt/utils/mmdc_ink.py @@ -1,51 +1,41 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- """ -@Time : 2023/7/4 10:53 -@Author : alexanderwu, imjohndoe +@Time : 2023/9/4 16:12 +@Author : alitrack @File : mermaid.py """ - -import requests import base64 import os -import click -@click.command() -@click.version_option() -@click.option("-i","--mermaid_code", type=str, help="mermaid code") -@click.option("-o","--output_file_without_suffix", type=str, help="output filename without suffix") -def mermaid_to_file(mermaid_code, output_file_without_suffix): +from aiohttp import ClientSession,ClientError +from metagpt.logs import logger + + +async def mermaid_to_file(mermaid_code, output_file_without_suffix): """suffix: png/svg :param mermaid_code: mermaid code :param output_file_without_suffix: output filename without suffix :return: 0 if succeed, -1 if failed """ - print('Starting mermaid_to_file command of mermaid.ink...') - encoded_string = base64.b64encode(mermaid_code.encode()).decode() - dir_name = os.path.dirname(output_file_without_suffix) - if dir_name and not os.path.exists(dir_name): - os.makedirs(dir_name) - with open(f"{output_file_without_suffix}.mmd", "w", encoding="utf-8") as f: - f.write(mermaid_code) - for suffix in ["svg", "png"]: output_file = f"{output_file_without_suffix}.{suffix}" path_type = "svg" if suffix == "svg" else "img" url = f"https://mermaid.ink/{path_type}/{encoded_string}" - response = requests.get(url) - - if response.status_code == 200: - with open(output_file, 'wb') as f: - f.write(response.content) - print(f"Generating {output_file}..") - else: - print(f"Failed to retrieve {suffix}") - return -1 - + async with ClientSession() as session: + try: + async with session.get(url) as response: + if response.status == 200: + text = await response.content.read() + with open(output_file, 'wb') as f: + f.write(text) + logger.info(f"Generating {output_file}..") + else: + logger.error(f"Failed to generate {output_file}") + return -1 + except ClientError as e: + logger.error(f"network error: {e}") + return -1 return 0 - -if __name__ == "__main__": - mermaid_to_file() \ No newline at end of file diff --git a/metagpt/utils/mmdc_playwright.py b/metagpt/utils/mmdc_playwright.py index d5d6b898e..bdbfd82ff 100644 --- a/metagpt/utils/mmdc_playwright.py +++ b/metagpt/utils/mmdc_playwright.py @@ -5,22 +5,13 @@ @Author : Steven Lee @File : mmdc_playwright.py """ + import os -import asyncio -import click from urllib.parse import urljoin - from playwright.async_api import async_playwright +from metagpt.logs import logger -@click.command() -@click.version_option() -@click.option("-i","--mermaid_code", type=str, help="mermaid code") -@click.option("-o","--output_file_without_suffix", type=str, help="output filename without suffix") -@click.option("--width",type=int,help="width",default=2048) -@click.option("--height",type=int,help="height",default=2048) -def mermaid_to_file(mermaid_code, output_file_without_suffix, width, height): - - +async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048)-> int: """ Converts the given Mermaid code to various output formats and saves them to files. @@ -33,104 +24,88 @@ def mermaid_to_file(mermaid_code, output_file_without_suffix, width, height): Returns: int: Returns 1 if the conversion and saving were successful, -1 otherwise. """ - + suffixes=['png', 'svg', 'pdf'] __dirname = os.path.dirname(os.path.abspath(__file__)) - async def mermaid_to_file0(mermaid_code, output_file_without_suffix, width=2048, height=2048, suffixes=['png', 'svg', 'pdf'])-> int: + async with async_playwright() as p: + browser = await p.chromium.launch() + device_scale_factor = 1.0 + context = await browser.new_context( + viewport={'width': width, 'height': height}, + device_scale_factor=device_scale_factor, + ) + page = await context.new_page() - async with async_playwright() as p: - browser = await p.chromium.launch() - device_scale_factor = 1.0 - context = await browser.new_context( - viewport={'width': width, 'height': height}, - device_scale_factor=device_scale_factor, - ) - page = await context.new_page() + async def console_message(msg): + logger.info(msg.text) + page.on('console', console_message) - async def console_message(msg): - print(msg.text) - page.on('console', console_message) + try: + await page.set_viewport_size({'width': width, 'height': height}) - try: - await page.set_viewport_size({'width': width, 'height': height}) + mermaid_html_path = os.path.abspath( + os.path.join(__dirname, 'index.html')) + mermaid_html_url = urljoin('file:', mermaid_html_path) + await page.goto(mermaid_html_url) + await page.wait_for_load_state("networkidle") - mermaid_html_path = os.path.abspath( - os.path.join(__dirname, 'index.html')) - mermaid_html_url = urljoin('file:', mermaid_html_path) - await page.goto(mermaid_html_url) - await page.wait_for_load_state("networkidle") + await page.wait_for_selector("div#container", state="attached") + mermaid_config = {} + background_color = "#ffffff" + my_css = "" + await page.evaluate(f'document.body.style.background = "{background_color}";') - await page.wait_for_selector("div#container", state="attached") - mermaid_config = {} - background_color = "#ffffff" - my_css = "" - await page.evaluate(f'document.body.style.background = "{background_color}";') + metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => { + const { mermaid, zenuml } = globalThis; + await mermaid.registerExternalDiagrams([zenuml]); + mermaid.initialize({ startOnLoad: false, ...mermaidConfig }); + const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container')); + document.getElementById('container').innerHTML = svg; + const svgElement = document.querySelector('svg'); + svgElement.style.backgroundColor = backgroundColor; - metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => { - const { mermaid, zenuml } = globalThis; - await mermaid.registerExternalDiagrams([zenuml]); - mermaid.initialize({ startOnLoad: false, ...mermaidConfig }); - const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container')); - document.getElementById('container').innerHTML = svg; - const svgElement = document.querySelector('svg'); - svgElement.style.backgroundColor = backgroundColor; + if (myCSS) { + const style = document.createElementNS('http://www.w3.org/2000/svg', 'style'); + style.appendChild(document.createTextNode(myCSS)); + svgElement.appendChild(style); + } - if (myCSS) { - const style = document.createElementNS('http://www.w3.org/2000/svg', 'style'); - style.appendChild(document.createTextNode(myCSS)); - svgElement.appendChild(style); - } + }''', [mermaid_code, mermaid_config, my_css, background_color]) - }''', [mermaid_code, mermaid_config, my_css, background_color]) - - if 'svg' in suffixes : - svg_xml = await page.evaluate('''() => { - const svg = document.querySelector('svg'); - const xmlSerializer = new XMLSerializer(); - return xmlSerializer.serializeToString(svg); - }''') - print(f"Generating {output_file_without_suffix}.svg..") - with open(f'{output_file_without_suffix}.svg', 'wb') as f: - f.write(svg_xml.encode('utf-8')) - - if 'png' in suffixes: - clip = await page.evaluate('''() => { - const svg = document.querySelector('svg'); - const rect = svg.getBoundingClientRect(); - return { - x: Math.floor(rect.left), - y: Math.floor(rect.top), - width: Math.ceil(rect.width), - height: Math.ceil(rect.height) - }; - }''') - await page.set_viewport_size({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height']}) - screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device') - print(f"Generating {output_file_without_suffix}.png..") - with open(f'{output_file_without_suffix}.png', 'wb') as f: - f.write(screenshot) - if 'pdf' in suffixes: - pdf_data = await page.pdf(scale=device_scale_factor) - print(f"Generating {output_file_without_suffix}.pdf..") - with open(f'{output_file_without_suffix}.pdf', 'wb') as f: - f.write(pdf_data) - return 1 - except Exception as e: - print(e) - return -1 - finally: - await browser.close() - dir_name = os.path.dirname(output_file_without_suffix) - if dir_name and not os.path.exists(dir_name): - os.makedirs(dir_name) - with open(f"{output_file_without_suffix}.mmd", "w", encoding="utf-8") as f: - f.write(mermaid_code) - suffixes = ['png', 'svg', 'pdf'] - loop = asyncio.new_event_loop() - result = loop.run_until_complete(mermaid_to_file0(mermaid_code, output_file_without_suffix, width, height, suffixes)) - loop.close() - return result - -if __name__ == "__main__": - mermaid_to_file() + if 'svg' in suffixes : + svg_xml = await page.evaluate('''() => { + const svg = document.querySelector('svg'); + const xmlSerializer = new XMLSerializer(); + return xmlSerializer.serializeToString(svg); + }''') + logger.info(f"Generating {output_file_without_suffix}.svg..") + with open(f'{output_file_without_suffix}.svg', 'wb') as f: + f.write(svg_xml.encode('utf-8')) + if 'png' in suffixes: + clip = await page.evaluate('''() => { + const svg = document.querySelector('svg'); + const rect = svg.getBoundingClientRect(); + return { + x: Math.floor(rect.left), + y: Math.floor(rect.top), + width: Math.ceil(rect.width), + height: Math.ceil(rect.height) + }; + }''') + await page.set_viewport_size({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height']}) + screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device') + logger.info(f"Generating {output_file_without_suffix}.png..") + with open(f'{output_file_without_suffix}.png', 'wb') as f: + f.write(screenshot) + if 'pdf' in suffixes: + pdf_data = await page.pdf(scale=device_scale_factor) + logger.info(f"Generating {output_file_without_suffix}.pdf..") + with open(f'{output_file_without_suffix}.pdf', 'wb') as f: + f.write(pdf_data) + return 0 + except Exception as e: + logger.error(e) + return -1 + finally: + await browser.close() diff --git a/metagpt/utils/mmdc_pyppeteer.py b/metagpt/utils/mmdc_pyppeteer.py index f3e00d053..56367236f 100644 --- a/metagpt/utils/mmdc_pyppeteer.py +++ b/metagpt/utils/mmdc_pyppeteer.py @@ -2,23 +2,15 @@ # -*- coding: utf-8 -*- """ @Time : 2023/9/4 16:12 -@Author : Steven Lee +@Author : alitrack @File : mmdc_pyppeteer.py """ -import asyncio -import click import os from urllib.parse import urljoin -import sys from pyppeteer import launch +from metagpt.logs import logger -@click.command() -@click.version_option() -@click.option("-i","--mermaid_code", type=str, help="mermaid code") -@click.option("-o","--output_file_without_suffix", type=str, help="output filename without suffix") -@click.option("--width",type=int,help="width",default=2048) -@click.option("--height",type=int,help="height",default=2048) -def mermaid_to_file(mermaid_code, output_file_without_suffix, width, height): +async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048)-> int: """ Converts the given Mermaid code to various output formats and saves them to files. @@ -31,104 +23,90 @@ def mermaid_to_file(mermaid_code, output_file_without_suffix, width, height): Returns: int: Returns 1 if the conversion and saving were successful, -1 otherwise. """ + suffixes = ['png', 'svg', 'pdf'] + __dirname = os.path.dirname(os.path.abspath(__file__)) - async def mermaid_to_file0(mermaid_code, output_file_without_suffix, width=2048, height=2048, suffixes=['png', 'svg', 'pdf'])-> int: - __dirname = os.path.dirname(os.path.abspath(__file__)) - executablePath = os.getenv('PUPPETEER_EXECUTABLE_PATH',"") - if executablePath: - browser = await launch(headless=True, - executablePath=executablePath, - args=['--disable-extensions',"--no-sandbox"] - ) - else: - print("Please set the environment variable:PUPPETEER_EXECUTABLE_PATH.") - return -1 - page = await browser.newPage() - device_scale_factor = 1.0 + executablePath = os.getenv('PUPPETEER_EXECUTABLE_PATH',"") + if executablePath: + browser = await launch(headless=True, + executablePath=executablePath, + args=['--disable-extensions',"--no-sandbox"] + ) + else: + logger.error("Please set the environment variable:PUPPETEER_EXECUTABLE_PATH.") + return -1 + page = await browser.newPage() + device_scale_factor = 1.0 - async def console_message(msg): - print(msg.text) - page.on('console', console_message) + async def console_message(msg): + logger.info(msg.text) + page.on('console', console_message) - try: - await page.setViewport(viewport={'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor}) + try: + await page.setViewport(viewport={'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor}) - mermaid_html_path = os.path.abspath( - os.path.join(__dirname, 'index.html')) - mermaid_html_url = urljoin('file:', mermaid_html_path) - await page.goto(mermaid_html_url) + mermaid_html_path = os.path.abspath( + os.path.join(__dirname, 'index.html')) + mermaid_html_url = urljoin('file:', mermaid_html_path) + await page.goto(mermaid_html_url) - await page.querySelector("div#container") - mermaid_config = {} - background_color = "#ffffff" - my_css = "" - await page.evaluate(f'document.body.style.background = "{background_color}";') + await page.querySelector("div#container") + mermaid_config = {} + background_color = "#ffffff" + my_css = "" + await page.evaluate(f'document.body.style.background = "{background_color}";') - metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => { - const { mermaid, zenuml } = globalThis; - await mermaid.registerExternalDiagrams([zenuml]); - mermaid.initialize({ startOnLoad: false, ...mermaidConfig }); - const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container')); - document.getElementById('container').innerHTML = svg; - const svgElement = document.querySelector('svg'); - svgElement.style.backgroundColor = backgroundColor; + metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => { + const { mermaid, zenuml } = globalThis; + await mermaid.registerExternalDiagrams([zenuml]); + mermaid.initialize({ startOnLoad: false, ...mermaidConfig }); + const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container')); + document.getElementById('container').innerHTML = svg; + const svgElement = document.querySelector('svg'); + svgElement.style.backgroundColor = backgroundColor; - if (myCSS) { - const style = document.createElementNS('http://www.w3.org/2000/svg', 'style'); - style.appendChild(document.createTextNode(myCSS)); - svgElement.appendChild(style); - } - }''', [mermaid_code, mermaid_config, my_css, background_color]) + if (myCSS) { + const style = document.createElementNS('http://www.w3.org/2000/svg', 'style'); + style.appendChild(document.createTextNode(myCSS)); + svgElement.appendChild(style); + } + }''', [mermaid_code, mermaid_config, my_css, background_color]) - if 'svg' in suffixes : - svg_xml = await page.evaluate('''() => { - const svg = document.querySelector('svg'); - const xmlSerializer = new XMLSerializer(); - return xmlSerializer.serializeToString(svg); - }''') - print(f"Generating {output_file_without_suffix}.svg..") - with open(f'{output_file_without_suffix}.svg', 'wb') as f: - f.write(svg_xml.encode('utf-8')) + if 'svg' in suffixes : + svg_xml = await page.evaluate('''() => { + const svg = document.querySelector('svg'); + const xmlSerializer = new XMLSerializer(); + return xmlSerializer.serializeToString(svg); + }''') + logger.info(f"Generating {output_file_without_suffix}.svg..") + with open(f'{output_file_without_suffix}.svg', 'wb') as f: + f.write(svg_xml.encode('utf-8')) - if 'png' in suffixes: - clip = await page.evaluate('''() => { - const svg = document.querySelector('svg'); - const rect = svg.getBoundingClientRect(); - return { - x: Math.floor(rect.left), - y: Math.floor(rect.top), - width: Math.ceil(rect.width), - height: Math.ceil(rect.height) - }; - }''') - await page.setViewport({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height'], 'deviceScaleFactor': device_scale_factor}) - screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device') - print(f"Generating {output_file_without_suffix}.png..") - with open(f'{output_file_without_suffix}.png', 'wb') as f: - f.write(screenshot) - if 'pdf' in suffixes: - pdf_data = await page.pdf(scale=device_scale_factor) - print(f"Generating {output_file_without_suffix}.pdf..") - with open(f'{output_file_without_suffix}.pdf', 'wb') as f: - f.write(pdf_data) - return 1 - except Exception as e: - print(e) - return -1 - finally: - await browser.close() - + if 'png' in suffixes: + clip = await page.evaluate('''() => { + const svg = document.querySelector('svg'); + const rect = svg.getBoundingClientRect(); + return { + x: Math.floor(rect.left), + y: Math.floor(rect.top), + width: Math.ceil(rect.width), + height: Math.ceil(rect.height) + }; + }''') + await page.setViewport({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height'], 'deviceScaleFactor': device_scale_factor}) + screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device') + logger.info(f"Generating {output_file_without_suffix}.png..") + with open(f'{output_file_without_suffix}.png', 'wb') as f: + f.write(screenshot) + if 'pdf' in suffixes: + pdf_data = await page.pdf(scale=device_scale_factor) + logger.info(f"Generating {output_file_without_suffix}.pdf..") + with open(f'{output_file_without_suffix}.pdf', 'wb') as f: + f.write(pdf_data) + return 0 + except Exception as e: + logger.error(e) + return -1 + finally: + await browser.close() - suffixes = ['png', 'svg', 'pdf'] - dir_name = os.path.dirname(output_file_without_suffix) - if dir_name and not os.path.exists(dir_name): - os.makedirs(dir_name) - with open(f"{output_file_without_suffix}.mmd", "w", encoding="utf-8") as f: - f.write(mermaid_code) - loop = asyncio.new_event_loop() - result = loop.run_until_complete(mermaid_to_file0(mermaid_code, output_file_without_suffix, width, height,suffixes)) - loop.close() - return result - -if __name__ == "__main__": - mermaid_to_file()