Add PNG support

Author: hkato

README.md

@@ -2,7 +2,7 @@
 
 [Mermaid](https://mermaid.js.org/) extension for [Python-Markdown](https://python-markdown.github.io/) using [mermaid-cli](https://github.com/mermaid-js/mermaid-cli).
 
-Mermaid code blocks are converted to SVG and treated as [data: URI](https://developer.mozilla.org/en-US/docs/Web/URI/Reference/Schemes/data). This allows for PDF generation with tools like [WeasyPrint](https://weasyprint.org/) without the need for JavaScript, even during web browsing.
+Mermaid code blocks are converted to SVG/PNG and treated as [data: URI](https://developer.mozilla.org/en-US/docs/Web/URI/Reference/Schemes/data). This allows for PDF generation with tools like [WeasyPrint](https://weasyprint.org/) without the need for JavaScript, even during web browsing.
 
 ## Install
 
@@ -24,14 +24,8 @@ markdown_text = """```mermaid
 sequenceDiagram
     participant Alice
     participant Bob
-    Alice->>John: Hello John, how are you?
-    loop HealthCheck
-        John->>John: Fight against hypochondria
-    end
-    Note right of John: Rational thoughts <br/>prevail!
-    John-->>Alice: Great!
-    John->>Bob: How about you?
-    Bob-->>John: Jolly good!
+    Bob->>Alice: Hi Alice
+    Alice->>Bob: Hi Bob
 ```"""
 
 html_output = markdown.markdown(markdown_text, extensions=[MermaidDataURIExtension()])
@@ -45,7 +39,7 @@ Gg6IDc1MHB4OyBiYWNrZ3JvdW5kLWNvbG9yOiB3aGl0ZTsiIHhtbG5zOnhsaW5rPSJodHRwOi8vd3d3L
 ...
 ...
 ...
-IHgxPSIyNzYiLz48L3N2Zz4=" alt="Mermaid diagram"></p>
+IHgxPSIyNzYiLz48L3N2Zz4=" ></p>
 ```
 
 ## MkDocs Integration
@@ -55,3 +49,22 @@ IHgxPSIyNzYiLz48L3N2Zz4=" alt="Mermaid diagram"></p>
 markdown_extensions:
   - mermaid_data_uri
 ```
+
+## Diagram
+
+```mermaid
+sequenceDiagram
+    participant application as Application<br/>(eg MkDocs)
+    participant markdown as Python Markdown
+    participant extension as MermaidDataURIExtension
+    participant mmdc as Mermaid CLI
+
+    application->>markdown: Markdown + Mermaid
+    markdown->>extension: Preprocessor
+    extension->>mmdc: Mermaid
+    mmdc-->>mmdc: Convert
+    mmdc-->>extension: Image Data
+    extension-->>extension: Base64 encode
+    extension-->>markdown: Markdown + data URI image
+    markdown-->>application: HTML + data URI image
+```

docs/mermaid.md

@@ -34,7 +34,7 @@ sequenceDiagram
 !!! warning
     When generating PDFs, text within the graph syntax doesn't get output in WeasyPrint. However, it does appear when using Playwright.
 
-```mermaid
+```mermaid image=png
 graph LR
     A[Square Rect] -- Link text --> B((Circle))
     A --> C(Round Rect)
@@ -47,7 +47,7 @@ graph LR
 !!! warning
     When generating PDFs, text within the graph syntax doesn't get output in WeasyPrint. However, it does appear when using Playwright.
 
-```mermaid
+```mermaid image=png
 graph TB
     sq[Square shape] --> ci((Circle shape))
 

src/mermaid_data_uri/extension.py

@@ -22,18 +22,42 @@ def run(self, lines: List[str]) -> List[str]:
         is_in_mermaid = False
 
         for line in lines:
-            if line.strip() == '```mermaid':
+            if line.strip().startswith('```mermaid'):
                 is_in_mermaid = True
                 mermaid_block: List[str] = []
+                # Extract options after '```mermaid'
+                options = line.strip()[10:].strip()
+                option_dict = {}
+                if options:
+                    for option in options.split():
+                        key, _, value = option.partition('=')
+                        option_dict[key] = value
                 continue
             elif line.strip() == '```' and is_in_mermaid:
                 is_in_mermaid = False
                 if mermaid_block:
                     mermaid_code = '\n'.join(mermaid_block)
-                    svg_content = self._mermaid2svg(mermaid_code)
-                    if svg_content:
-                        data_uri = self._svg2data_uri(svg_content)
-                        new_lines.append(f'<img src="{data_uri}" alt="Mermaid diagram">')
+
+                    # Image type handling
+                    if 'image' in option_dict:
+                        image_type = option_dict['image']
+                        del option_dict['image']
+                        if image_type not in ['svg', 'png']:
+                            image_type = 'svg'
+                    else:
+                        image_type = 'svg'
+
+                    base64image = self._mermaid2base64image(mermaid_code, image_type)
+                    if base64image:
+                        # Build the <img> tag with extracted options
+                        if image_type == 'svg':
+                            img_tag = f'<img src="data:image/svg+xml;base64,{base64image}"'
+                        else:
+                            img_tag = f'<img src="data:image/png;base64,{base64image}"'
+                        for key, value in option_dict.items():
+                            img_tag += f' {key}={value}'
+                        img_tag += ' />'
+                        new_lines.append(img_tag)
                     else:
                         new_lines.append('```mermaid')
                         new_lines.extend(mermaid_block)
@@ -47,41 +71,47 @@ def run(self, lines: List[str]) -> List[str]:
 
         return new_lines
 
-    def _svg2data_uri(self, svg_content: str) -> str:
-        """Convert SVG content to Data URI."""
-        base64_svg = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
-        data_uri = f'data:image/svg+xml;base64,{base64_svg}'
-        return data_uri
-
-    def _mermaid2svg(self, mermaid_code: str) -> str:
+    def _mermaid2base64image(self, mermaid_code: str, image_type: str) -> str:
         """Convert mermaid code to SVG using mmdc (Mermaid CLI)."""
         with tempfile.NamedTemporaryFile(mode='w', suffix='.mmd', delete=False) as tmp_mmd:
             tmp_mmd.write(mermaid_code)
             mmd_filepath = tmp_mmd.name
 
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.svg', delete=False) as tmp_svg:
-            svg_filepath = tmp_svg.name
+        with tempfile.NamedTemporaryFile(mode='w', suffix=f'.{image_type}', delete=False) as tmp_img:
+            img_filepath = tmp_img.name
 
         try:
             command = [
                 'mmdc',
                 '--input',
                 mmd_filepath,
                 '--output',
-                svg_filepath,
+                img_filepath,
+                '--outputFormat',
+                image_type,
                 '--puppeteerConfigFile',
                 os.path.join(os.path.dirname(__file__), 'puppeteer-config.json'),
             ]
             subprocess.run(command, check=True, capture_output=True)
-            with open(svg_filepath, 'r', encoding='utf-8') as f:
-                svg_content = f.read()
-            return svg_content
+            if image_type == 'svg':
+                with open(img_filepath, 'r', encoding='utf-8') as f:
+                    svg_content: str = f.read()
+            elif image_type == 'png':
+                with open(img_filepath, 'rb') as f:
+                    png_content: bytes = f.read()
         except subprocess.CalledProcessError as e:
             print(f'Error generating SVG: {e.stderr.decode()}')
             return ''
         finally:
             os.remove(mmd_filepath)
-            os.remove(svg_filepath)
+            os.remove(img_filepath)
+
+        if image_type == 'svg':
+            base64image = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
+        else:
+            base64image = base64.b64encode(png_content).decode('utf-8')
+
+        return base64image
 
 
 class MermaidDataURIExtension(Extension):