My Custom Code Block Setup in Astro

I wanted to make the code blocks on this site a little better. When I write a post with a code snippet, I want to show the filename, add line numbers, and sometimes even start the numbering from a specific line.

I was hoping I could just write my markdown like this:

Markdown

```rust file=src/main.rs start=12 numbers
// Some rust code
```

And have it automatically generate a nice header with the filename and handle the line numbering.

My site is built with Astro, and it uses Shiki for syntax highlighting. After digging into the Astro docs, I found out you can customize how Shiki works using something called a Shiki Transformer. This sounded like exactly what I needed.

Shiki Transformer

A Shiki transformer lets you hook into the process that turns your code block into highlighted HTML. You can inspect the code, look at the metadata you passed in (like file=src/main.rs), and change the final HTML pre tag.

So, I wrote a transformer to do three things:

Grab the filename from the file attribute.
Add a CSS class if numbers is present.
Add a CSS variable for the start line number.

Here’s the transformer I came up with:

src/plugins/shiki-code-metadata.ts

import type { ShikiTransformer } from "shiki";

// A map to add default filenames if I forget to specify one
const langToFilenameMap: Record<string, string> = {
    bash: "Shell",
    c: "C",
    css: "CSS",
    go: "Go",
    html: "HTML",
    js: "JavaScript",
    json: "JSON",
    plaintext: "Plain Text",
    py: "Python",
    python: "Python",
    ruby: "Ruby",
    rust: "Rust",
    scala: "Scala",
    sh: "Shell",
    shell: "Shell",
    solidity: "Solidity",
    toml: "TOML",
    ts: "TypeScript",
    tsx: "TSX",
    typescript: "TypeScript",
    yaml: "YAML",
    zsh: "Shell"
};

export default function shikiCodeMetadata(): ShikiTransformer {
    return {
        pre(node) {
            // The metadata string is available in this.options.meta.__raw
            const meta = this.options.meta?.__raw;

            // If there's no metadata, we might still want a default filename
            if (meta == undefined) {
                const lang = this.options.lang;
                const defaultFilename = langToFilenameMap[lang];
                if (defaultFilename) {
                    node.properties.dataFilename = defaultFilename;
                }
                return;
            }

            // If we have metadata, let's parse it
            const attributes = meta.split(/\s+/).filter(Boolean);
            const styles: string[] = [];
            let hasExplicitFilename = false;

            for (const attr of attributes) {
                const [key, value] = attr.split("=", 2);

                if (!key) {
                    continue;
                }

                switch (key) {
                    case "file":
                        node.properties.dataFilename = value;
                        hasExplicitFilename = true;
                        break;
                    case "start":
                        // This sets a CSS variable on the <pre> element
                        styles.push(`--start-line: ${value}`);
                        break;
                    case "numbers":
                        // This just adds a class for styling with CSS
                        this.addClassToHast(node, "show-line-numbers");
                        break;
                }
            }

            // Add a default filename if one wasn't provided
            if (!hasExplicitFilename) {
                const lang = this.options.lang;
                const defaultFilename = langToFilenameMap[lang];
                if (defaultFilename) {
                    node.properties.dataFilename = defaultFilename;
                }
            }

            // Combine any new styles with existing ones
            if (styles.length > 0) {
                const existingStyle = (node.properties.style as string) || "";
                const separator =
                    existingStyle && !existingStyle.endsWith(";") ? ";" : "";
                node.properties.style =
                    existingStyle + separator + styles.join(";");
            }
        }
    };
}

This transformer attaches the filename to the <pre> tag as a data-filename attribute. I can use this attribute later to build the UI around the code block. The numbers flag adds a show-line-numbers class, which I can use in my CSS to show the line counter.

The Rehype Plugin

The transformer modified the <pre> tag, but I still needed to add the header with the filename and a copy button. I couldn’t do this inside the Shiki transformer because it only has access to the <pre> and <code> elements.

This is where I learned about rehype. After Shiki generates the HTML for the code block, Astro can run rehype plugins to modify that HTML. This was the perfect place to add my wrapper.

My rehype plugin looks for any <pre> tag that has the data-filename attribute my Shiki transformer added. If it finds one, it creates a new div to act as a container, builds the header, and then moves the original <pre> tag inside it.

Here is the rehype plugin:

src/plugins/rehype-code-wrapper.ts

import type { Root } from "hast";
import { visit } from "unist-util-visit";
import { h } from "hastscript";

export default function rehypeCodeWrapper() {
    return function (tree: Root) {
        visit(tree, "element", function (node, index, parent) {
            // Only act on <pre> elements
            if (node.tagName != "pre") {
                return;
            }

            // Check for the data attribute we added in the Shiki transformer
            if (!node.properties || !node.properties.dataFilename) {
                return;
            }

            const filename = node.properties.dataFilename as string;

            // Create the new structure
            const wrapper = h("div.code-container", [
                h("div.code-header", [
                    h("div.code-filename", [filename]),
                    h("button.code-copy", ["Copy"])
                ]),
                node // The original <pre> element
            ]);

            // Replace the original node with our new wrapper
            if (parent && index != null) {
                parent.children[index] = wrapper;
            }
        });
    };
}

Plugging It Into Astro

The last step was to tell Astro to use my new plugins. I just had to add them to my astro.config.mjs file.

astro.config.mjs

import rehypeCodeWrapper from "./src/plugins/rehype-code-wrapper";
import shikiCodeMetadata from "./src/plugins/shiki-code-metadata";

export default defineConfig({
    markdown: {
        shikiConfig: {
            transformers: [
                // Our Shiki transformer runs first
                shikiCodeMetadata()
            ]
        },
        rehypePlugins: [
            // Then our rehype plugin runs on the HTML
            rehypeCodeWrapper
        ]
    }
});

Now, every code block I write in markdown gets processed by these two plugins, and the final HTML looks like this:

HTML

<div class="code-container">
    <div class="code-header">
        <div class="code-filename">src/main.rs</div>
        <button class="code-copy">Copy</button>
    </div>
    <pre class="show-line-numbers" style="--start-line: 12">
		    <!-- highlighted code here -->
    </pre>
</div>

This was a fun little coding exercise. The main thing I learned was the difference between remark and rehype. Remark work with the markdown structure before it becomes HTML, which is what Shiki does. Rehype work on the HTML structure after it has been generated. Understanding that made it clear how to solve this problem.

Shiki Transformer

The Rehype Plugin

Plugging It Into Astro

Topics