Learn Shiki Rehype Pretty Code - Custom HAST & Advanced Rehype Plugins

#Introduction

rehype-pretty-code produces complete markup, but real needs often go beyond its built-in features: adding a language to the copy button, changing certain classes, or injecting metadata from outside. All of that can be done by modifying the HAST tree after the plugin finishes.

Episode 17 covers custom HAST and advanced rehype plugins: understanding the shape of HAST nodes, writing custom rehype plugins that visit and modify nodes, and arranging the correct plugin order in complex pipelines.

#Understanding HAST

#The Node Tree for HTML

HAST is the tree representation of HTML used by the rehype pipeline. Elements like figure, pre, and code are element nodes with tagName, properties, and children properties. Attributes like data-rehype-pretty-code-figure are stored in properties as key-value pairs.

Here's the shape of one of rehype-pretty-code's output nodes:

The shape of a figure node

{
  type: "element",
  tagName: "figure",
  properties: { dataRehypePrettyCodeFigure: "" },
  children: [{ type: "element", tagName: "pre", children: [] }]
}

Understanding this structure is the basis for writing plugins that target specific elements. All your modifications end up as attributes on the final HTML.

#Why Modify HAST

Modifying the tree is safer than manipulating HTML strings. Characters that need escaping are still handled by rehype-stringify, and attributes added through properties are turned into valid HTML attributes. This keeps the output clean and avoids escaping bugs.

#Creating a Custom Rehype Plugin

#Plugin Structure

A rehype plugin is a function that accepts options and returns a transformer. The transformer receives the HAST tree and may modify it directly:

Plugin to add a language attribute

import { visit } from "unist-util-visit";
 
export function rehypeAddLangAttribute(options) {
  return (tree) => {
    visit(tree, { tagName: "figure" }, (node) => {
      const code = node.children.find(
        (child) => child.tagName === "code",
      );
      const lang = code?.properties?.dataLanguage;
      if (lang) {
        node.properties.dataLang = lang;
      }
    });
  };
}

The plugin above uses unist-util-visit to find all figure elements. Every figure containing a code with a language gets a data-lang attribute readable by client-side JavaScript.

#Installing the Plugin in the Pipeline

Register the plugin after rehype-pretty-code so the nodes are already in figure form when accessed:

Installation order

unified()
  .use(remarkParse)
  .use(remarkRehype)
  .use(rehypePrettyCode, { theme: "github-dark-default" })
  .use(rehypeAddLangAttribute)
  .use(rehypeStringify);

This order guarantees the custom plugin sees a complete structure. If installed before rehype-pretty-code, no figure nodes are found and the plugin does nothing.

#Modifying Figure and Pre

#Adding Attributes to Pre

Sometimes you need to mark <pre> so CSS can target it specifically. Add the attribute directly to the node's properties:

Add an attribute to pre

import { visit } from "unist-util-visit";
 
export function rehypeMarkCodeBlocks() {
  return (tree) => {
    visit(tree, { tagName: "pre" }, (node) => {
      node.properties.dataCodeBlock = "";
    });
  };
}

After this plugin runs, every <pre> carries data-code-block. CSS and client JavaScript can use it without relying on a figure structure that may differ between plugin versions.

#Inserting New Nodes

A plugin can also add elements, for example wrapping the code with a wrapper for styling. Use a regular element node and place it in children:

Insert a wrapper

visit(tree, { tagName: "figure" }, (node) => {
  node.children.push({
    type: "element",
    tagName: "div",
    properties: { className: ["code-footer"] },
    children: [],
  });
});

Inserted nodes must follow the HAST format: type, tagName, properties, and children. rehype-stringify will turn them into correct HTML.

#Plugin Order in the Pipeline

#A Common Chain

A production pipeline often combines many plugins. An example of a sensible order:

Complete pipeline

unified()
  .use(remarkParse)
  .use(remarkRehype)
  .use(rehypeRaw)
  .use(rehypePrettyCode, { theme: "github-dark-default" })
  .use(rehypeSlug)
  .use(rehypeAutolinkHeadings)
  .use(rehypeSanitize)
  .use(rehypeStringify);

rehypeRaw turns raw HTML inside Markdown into nodes, then rehypePrettyCode processes code blocks, followed by rehypeSlug and rehypeAutolinkHeadings for headings, and sanitization at the end.

#Rules for Determining Order

The principle is simple: plugins that change the structure must run before plugins that read that structure. Languages are highlighted before headings get ids, and sanitization is always at the end to clean up all previous output. Swapping the order can make a plugin lose its target or let unsafe markup through.

#Conclusion

Key takeaways:

• HAST is a tree of element nodes with tagName, properties, and children.
• Modifying HAST is safer than manipulating HTML strings.
• Custom rehype plugins are made from functions that return a transformer.
• unist-util-visit helps find specific nodes in the tree.
• New nodes are inserted in complete HAST format.
• Plugin order determines whether transformations support or break each other.

In episode 18 you'll learn Shiki core and the highlighter API: using createHighlighter, createHighlighterCore, codeToHtml, codeToTokens, and getSingletonHighlighter, plus using Shiki outside rehype like markdown-it, a CLI, or manual rendering.