Adding UML Diagrams to Rehype in Astro

October 7, 2024 David Teather

I just recently added UML diagrams for my blogs, since in the future I want to create blog posts and youtube videos about system design and wanted to define UML diagrams in markdown and have them rendered out.

In this blog post I’ll show you how I got UML diagrams like these to work with my current setup of Rehype + Astro.

I also wanted to be able to put images in which was annoying enough to encourage me to make a blog post about this

I decided to use mermaidjs to define and render out the diagrams since it seems to get the job done and honestly it had a better website and more stars than alternatives on GitHub. Let’s get started!

Initial Setup

We’ll be using rehype-mermaid, the way this library and others work is it internally uses playwright which is a headless browser that uses the mermaid javascript library to generate the images. Unfortunantely this does make both installation and the build step need an additional dependency.

Lets install this package with

npm install rehype-mermaid

And install playwright’s browsers with

npx playwright-core install --with-deps chromium

Then lets add rehype mermaid to our astro config, if you’re new to astro maybe checkout the official rehype mermaid example

astro.config.mjs

import rehypeMermaid from 'rehype-mermaid'
 
export default defineConfig({
    markdown: {
        //  add to your existing array plugins
        rehypePlugins: [..., rehypeMermaid, ...]
    },
});

Then after that, you should be pretty much all set for the initial setup, try to create a mermaid codeblock on a markdown file of yours. Here’s an example from earlier

```mermaid
graph TD;
    A-->B;
    A-->C;
    B-->D;
    C-->D;
```

If that works for you great! However, I had some things that I wanted to change

It was getting cut off (because of other styles that were getting applied)
It always centered left, I wanted it to center at the middle
I wanted to support dark mode

Configuring and Styling

So let’s address these! I’ll be creating a post-processor rehype extension to do these, it’s probably not the most efficient or best practice way since I don’t know rehype that well but hey it works! It also just runs at build time so it won’t slow users down.

Fix The Cutoff

Some of my styling was overlapping and it wasn’t rendering as expected. Luckily our rehype mermaid plugin has a few different modes it can work in.

inline-svg (default)
- directly embeds the svg mermaid generates into the page inline
img-png
- playwright takes a png and embeds that into the page
img-svg
- embeds the svg as a <img> into the page
pre-mermaid
- renders mermaid on the client

What solved this issue for me was using the img-svg. I still get the benefits of svg and running mermaid at build time, but the styling doesn’t get applied.

Update the strategy to use img-svg

astro.config.mjs

import rehypeMermaid from 'rehype-mermaid'
 
export default defineConfig({
    markdown: {
        rehypePlugins: [..., [rehypeMermaid, {
            strategy: 'img-svg',
        }], ...]
    },
});

Hopefully now it’s not getting cut off anymore!

Center Middle

To center the diagrams in the middle I added a post-processor rehype plugin. Here’s my code for it, all it does is add an mx-auto since I’m using tailwind.

function rehypeModifyMermaidGraphs() {
    return (tree) => {
        visit(tree, 'element', (node, index, parent) => {
            if (node.tagName === 'img') {
                // check the prefix starts with mermaid
                const split = node.properties.id.split('-');
                if (split.length !== 2) return;
                if (split[0] !== 'mermaid') return; 
 
                // add mx-auto
                node.properties.className = node.properties.className || [];
                node.properties.className.push('mx-auto');
            }
        });
    };
}

Then add it after mermaid in your rehype plugins

astro.config.mjs

import rehypeMermaid from 'rehype-mermaid'
 
export default defineConfig({
    markdown: {
        rehypePlugins: [..., [rehypeMermaid, {
            strategy: 'img-svg',
        }], rehypeModifyMermaidGraphs, ...]
    },
});

Now that should center your diagrams.

Dark Mode

This is another option we can provide to rehype mermaid, however it does require that we’re using either the img-png or img-svg strategies which we already are from earlier!

Update our configuration one last time

astro.config.mjs

import rehypeMermaid from 'rehype-mermaid'
 
export default defineConfig({
    markdown: {
        rehypePlugins: [..., [rehypeMermaid, {
            strategy: 'img-svg',
            dark: true
        }], rehypeModifyMermaidGraphs, ...]
    },
});

This will automatically generate an html structure like

<picture>
    <source srcset="darkmode_data" media="(prefers-color-scheme: dark)">
    <img src="lightmode_data">
</picture>

If you’re not familiar with what <source>, it switches out your <img> when the media condition applies. I also didn’t know about it before doing this :D

This should work if your website is set up to do dark mode based on the media query alone, however I’m using a manual toggle which applies data-theme="dark" to the root level html. This means I had to do it a little bit more manual and wanted to flip between media="all" and media="none" depending on what theme the user had selected.

For this I extended my mermaid post-processor to support the new dark: true setting and added classes mermaid and mermaid-dark to <img> and <source> respectively. Most of the code change is just conditional checking

function rehypeModifyMermaidGraphs() {
    return (tree) => {
        visit(tree, 'element', (node, index, parent) => {
            if (node.tagName === 'picture') {
                // if has children <source> and <img> with a tag prefix of "mermaid-[index]"
                if (node.children.length === 2 && node.children[0].tagName === 'source' && node.children[1].tagName === 'img') {
                    const sourceNode = node.children[0];
                    const imgNode = node.children[1];
 
                    // check the prefix
                    const split = imgNode.properties.id.split('-');
                    if (split.length !== 2) return;
                    if (split[0] !== 'mermaid') return; 
 
                    const mermaidIndex = parseInt(split[1]);
                    const darkID = `mermaid-dark-${mermaidIndex}`;
                    const ID = `mermaid-${mermaidIndex}`;
 
                    // skip if the children IDs don't match
                    if (sourceNode.properties.id !== darkID || imgNode.properties.id !== ID) {
                        return;
                    }
 
                    // add mx-auto to both source and img
                    sourceNode.properties.className = sourceNode.properties.className || [];
                    sourceNode.properties.className.push('mx-auto');
                    sourceNode.properties.className.push('mermaid-dark');
 
                    imgNode.properties.className = imgNode.properties.className || [];
                    imgNode.properties.className.push('mx-auto');
                    imgNode.properties.className.push('mermaid');
                }
            }
        });
    };
}

Then I added the following code in my app that gets triggered after the user has clicked the toggle theme button.

const dataTheme = document.documentElement.getAttribute("data-theme");
document.querySelectorAll(".mermaid-dark").forEach((el) => {
    if (dataTheme === "dark") {
    el.setAttribute("media", "all");
    } else {
    el.setAttribute("media", "none");
    }
});

If you’re saving the state of darkmode or lightmode between user refreshes, you’ll want to use a<script is:inline> to run the same code after you know your diagrams have been generated to avoid flashing from light to dark. I put the inline script at the bottom of my article layout.

Now your diagrams should look like what you see on this page! Test them out with the theme toggle in the top as well.

Adding Image Support

Ok so here’s where most of the hacky solutions come in for generating the diagrams. If you have better solutions please let me know.

Notes:

Hacky, doesn’t cover all svg edge cases
- I’m minifying on the svgs which seems to help reduce the total edge cases
Ideally just use the native-svg option but I couldn’t since I had some conflicting styling and wanted to use the default dark mode
Only works with svgs
I didn’t want to just link to a public host since I wanted to develop it locally and test it

I wanted to be able to support images like the following.

The code for this is this, this only supports svgs and the icons folder is at the astro public/icons folder.

```mermaid
 flowchart LR
     A[<img src='/icons/dead-bird.svg' width='50' height='50' /> ]
     A --> B[<img src='/icons/worm.svg' width='50' height='50' /> ]
```

Unfortunantely I solved this with yet another rehype extension, but this time it’s a preprocessor that will read the svg file and inline the contents of it into the mermaid code.

It’s pretty messy but here it is

function rehypeInlineSvg() {
    return (tree) => {
        visit(tree, 'element', (node) => {
            if (node.tagName === 'code' && node.properties.className && node.properties.className.length > 0 && node.properties.className[0] === 'language-mermaid') {
                // We need to look up any elements in the value that are like <img src='/icons/worm.svg' width='25' height='25' />
                // and then directly replace the img with the svg content inline
                var mermaidCode = node.children[0].value;
 
                // regex to find all the img tags
                const imgRegex = /<img[^>]*\/>/g;
                const imgMatches = mermaidCode.match(imgRegex);
 
                if (imgMatches) {
                    for (const imgMatch of imgMatches) {
                        const imgNode = unified().use(rehypeParse).parse(imgMatch);
                        const imgElement = imgNode.children[0].children[1].children[0];
                        const src = imgElement.properties.src;
 
                        // Check if the file is an SVG if so inline time
                        if (src.endsWith('.svg')) {
                            let svgPath;
 
                            if (src.startsWith('/')) {
                                const publicFolder = process.env.PUBLIC_FOLDER_PATH || 'public';
                                svgPath = resolve(publicFolder, `.${src}`);
                            } else {
                                svgPath = join(process.cwd(), src);
                            }
 
                            try {
                                var svgContent = readFileSync(svgPath, 'utf-8');
 
                                // cut off anything before the first <svg tag
                                svgContent = svgContent.substring(svgContent.indexOf('<svg'));
 
                                const inlineSvgNode = unified()
                                    .use(rehypeParse, { fragment: true })
                                    .parse(svgContent);
 
                                const svgElement = inlineSvgNode.children[0];
 
                                // Modify each <path> element inside the <svg>
                                // You might want to modify this to be conditional
                                visit(svgElement, 'element', (svgNode) => {
                                    if (svgNode.tagName === 'path') {
                                        svgNode.properties.style = svgNode.properties.style || '';
 
                                        var fill = svgNode.properties.fill;
                                        svgNode.properties.style += `fill: ${fill || 'inherit'} !important;`;
                                        svgNode.properties.style += `stroke-width: 0 !important;`;
                                        delete svgNode.properties.fill;
                                    }
                                });
 
                                // Copy relevant properties from the <img> to the <svg> tag
                                svgElement.properties = svgElement.properties || {};
 
                                if (imgElement.properties.width) {
                                    svgElement.properties.width = imgElement.properties.width;
                                }
                                if (imgElement.properties.height) {
                                    svgElement.properties.height = imgElement.properties.height;
                                }
                                if (imgElement.properties.style) {
                                    svgElement.properties.style = imgElement.properties.style;
                                }
                                if (imgElement.properties.className) {
                                    svgElement.properties.className = imgElement.properties.className;
                                }
 
                                const inlineSvgHtml = unified().use(rehypeStringify).stringify(inlineSvgNode).replaceAll("\n", "").replaceAll("\t", "");
 
                                // Replace with the inline SVG code
                                mermaidCode = mermaidCode.replace(imgMatch, inlineSvgHtml);
                            } catch (error) {
                                console.error(`Failed to inline SVG at ${svgPath}:`, error);
                            }
                        }
                    }
                    node.children[0].value = mermaidCode;
                }
            }
        });
    };
}

It could be worse, but it’s pretty sloppy. Then add it before our mermaid rehype plugin

astro.config.mjs

import rehypeMermaid from 'rehype-mermaid'
 
export default defineConfig({
    markdown: {
        rehypePlugins: [..., rehypeInlineSvg, [rehypeMermaid, {
            strategy: 'img-svg',
            dark: true
        }], rehypeModifyMermaidGraphs, ...]
    },
});

Updating Deployment

Ok finally since we used playwright we need to make sure to add the following to our build steps before npm build

npm install -D @playwright/test@latest
npx playwright install --with-deps

Now hopefully everything gets built properly.

I hope this helped you out and if you’re interested in any specific system design questions or blog posts definitely let me know :D

Have a great day ❤️

Back to blog