Refactor createShikiHighlighter (#11825)

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

.changeset/breezy-colts-promise.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/markdoc': patch
+---
+
+Uses latest version of `@astrojs/markdown-remark` with updated Shiki APIs

.changeset/hungry-jokes-try.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/markdown-remark': major
+---
+
+Updates return object of `createShikiHighlighter` as `codeToHast` and `codeToHtml` to allow generating either the hast or html string directly

.changeset/large-zebras-sniff.md

@@ -0,0 +1,9 @@
+---
+'astro': major
+---
+
+Updates internal Shiki rehype plugin to highlight code blocks as hast (using Shiki's `codeToHast()` API). This allows a more direct Markdown and MDX processing, and improves the performance when building the project, but may cause issues with existing Shiki transformers.
+
+If you are using Shiki transformers passed to `markdown.shikiConfig.transformers`, you must make sure they do not use the `postprocess` hook as it no longer runs on code blocks in `.md` and `.mdx` files. (See [the Shiki documentation on transformer hooks](https://shiki.style/guide/transformers#transformer-hooks) for more information). 
+
+Code blocks in `.mdoc` files and `<Code />` component do not use the internal Shiki rehype plugin and are unaffected.

packages/astro/components/Code.astro

@@ -111,13 +111,13 @@ const highlighter = await getCachedHighlighter({
 	],
 	theme,
 	themes,
-	defaultColor,
-	wrap,
-	transformers,
 });
 
-const html = await highlighter.highlight(code, typeof lang === 'string' ? lang : lang.name, {
+const html = await highlighter.codeToHtml(code, typeof lang === 'string' ? lang : lang.name, {
+	defaultColor,
+	wrap,
 	inline,
+	transformers,
 	meta,
 	attributes: rest as any,
 });

packages/integrations/markdoc/src/extensions/shiki.ts

@@ -5,7 +5,11 @@ import { unescapeHTML } from 'astro/runtime/server/index.js';
 import type { AstroMarkdocConfig } from '../config.js';
 
 export default async function shiki(config?: ShikiConfig): Promise<AstroMarkdocConfig> {
-	const highlighter = await createShikiHighlighter(config);
+	const highlighter = await createShikiHighlighter({
+		langs: config?.langs,
+		theme: config?.theme,
+		themes: config?.themes,
+	});
 
 	return {
 		nodes: {
@@ -16,7 +20,11 @@ export default async function shiki(config?: ShikiConfig): Promise<AstroMarkdocC
 					// Only the `js` part is parsed as `attributes.language` and the rest is ignored. This means
 					// some Shiki transformers may not work correctly as it relies on the `meta`.
 					const lang = typeof attributes.language === 'string' ? attributes.language : 'plaintext';
-					const html = await highlighter.highlight(attributes.content, lang);
+					const html = await highlighter.codeToHtml(attributes.content, lang, {
+						wrap: config?.wrap,
+						defaultColor: config?.defaultColor,
+						transformers: config?.transformers,
+					});
 
 					// Use `unescapeHTML` to return `HTMLString` for Astro renderer to inline as HTML
 					return unescapeHTML(html) as any;

packages/markdown/remark/src/highlight.ts

@@ -4,7 +4,11 @@ import { toText } from 'hast-util-to-text';
 import { removePosition } from 'unist-util-remove-position';
 import { visitParents } from 'unist-util-visit-parents';
 
-type Highlighter = (code: string, language: string, options?: { meta?: string }) => Promise<string>;
+type Highlighter = (
+	code: string,
+	language: string,
+	options?: { meta?: string },
+) => Promise<Root | string>;
 
 const languagePattern = /\blanguage-(\S+)\b/;
 
@@ -73,12 +77,17 @@ export async function highlightCodeBlocks(tree: Root, highlighter: Highlighter)
 	for (const { node, language, grandParent, parent } of nodes) {
 		const meta = (node.data as any)?.meta ?? node.properties.metastring ?? undefined;
 		const code = toText(node, { whitespace: 'pre' });
-		// TODO: In Astro 5, have `highlighter()` return hast directly to skip expensive HTML parsing and serialization.
+		const result = await highlighter(code, language, { meta });
-		const html = await highlighter(code, language, { meta });
+
-		// The replacement returns a root node with 1 child, the `<pr>` element replacement.
+		let replacement: Element;
-		const replacement = fromHtml(html, { fragment: true }).children[0] as Element;
+		if (typeof result === 'string') {
+			// The replacement returns a root node with 1 child, the `<pre>` element replacement.
+			replacement = fromHtml(result, { fragment: true }).children[0] as Element;
 			// We just generated this node, so any positional information is invalid.
 			removePosition(replacement);
+		} else {
+			replacement = result.children[0] as Element;
+		}
 
 		// We replace the parent in its parent with the new `<pre>` element.
 		const index = grandParent.children.indexOf(parent);

packages/markdown/remark/src/index.ts

@@ -26,7 +26,12 @@ export { rehypeHeadingIds } from './rehype-collect-headings.js';
 export { remarkCollectImages } from './remark-collect-images.js';
 export { rehypePrism } from './rehype-prism.js';
 export { rehypeShiki } from './rehype-shiki.js';
-export { createShikiHighlighter, type ShikiHighlighter } from './shiki.js';
+export {
+	createShikiHighlighter,
+	type ShikiHighlighter,
+	type CreateShikiHighlighterOptions,
+	type ShikiHighlighterHighlightOptions,
+} from './shiki.js';
 export * from './types.js';
 
 export const markdownConfigDefaults: Required<AstroMarkdownOptions> = {

packages/markdown/remark/src/rehype-shiki.ts

@@ -8,9 +8,20 @@ export const rehypeShiki: Plugin<[ShikiConfig?], Root> = (config) => {
 	let highlighterAsync: Promise<ShikiHighlighter> | undefined;
 
 	return async (tree) => {
-		highlighterAsync ??= createShikiHighlighter(config);
+		highlighterAsync ??= createShikiHighlighter({
+			langs: config?.langs,
+			theme: config?.theme,
+			themes: config?.themes,
+		});
 		const highlighter = await highlighterAsync;
 
-		await highlightCodeBlocks(tree, highlighter.highlight);
+		await highlightCodeBlocks(tree, (code, language, options) => {
+			return highlighter.codeToHast(code, language, {
+				meta: options?.meta,
+				wrap: config?.wrap,
+				defaultColor: config?.defaultColor,
+				transformers: config?.transformers,
+			});
+		});
 	};
 };

packages/markdown/remark/src/shiki.ts

@@ -1,25 +1,63 @@
-import type { Properties } from 'hast';
+import type { Properties, Root } from 'hast';
 import {
 	type BundledLanguage,
+	type LanguageRegistration,
+	type ShikiTransformer,
+	type ThemeRegistration,
+	type ThemeRegistrationRaw,
 	createCssVariablesTheme,
-	getHighlighter,
+	createHighlighter,
 	isSpecialLang,
 } from 'shiki';
-import type { ShikiConfig } from './types.js';
+import type { ThemePresets } from './types.js';
 
 export interface ShikiHighlighter {
-	highlight(
+	codeToHast(
 		code: string,
 		lang?: string,
-		options?: {
+		options?: ShikiHighlighterHighlightOptions,
+	): Promise<Root>;
+	codeToHtml(
+		code: string,
+		lang?: string,
+		options?: ShikiHighlighterHighlightOptions,
+	): Promise<string>;
+}
+
+export interface CreateShikiHighlighterOptions {
+	langs?: LanguageRegistration[];
+	theme?: ThemePresets | ThemeRegistration | ThemeRegistrationRaw;
+	themes?: Record<string, ThemePresets | ThemeRegistration | ThemeRegistrationRaw>;
+}
+
+export interface ShikiHighlighterHighlightOptions {
+	/**
+	 * Generate inline code element only, without the pre element wrapper.
+	 */
 	inline?: boolean;
+	/**
+	 * Enable word wrapping.
+	 * - true: enabled.
+	 * - false: disabled.
+	 * - null: All overflow styling removed. Code will overflow the element by default.
+	 */
+	wrap?: boolean | null;
+	/**
+	 * Chooses a theme from the "themes" option that you've defined as the default styling theme.
+	 */
+	defaultColor?: 'light' | 'dark' | string | false;
+	/**
+	 * Shiki transformers to customize the generated HTML by manipulating the hast tree.
+	 */
+	transformers?: ShikiTransformer[];
+	/**
+	 * Additional attributes to be added to the root code block element.
+	 */
 	attributes?: Record<string, string>;
 	/**
-			 * Raw `meta` information to be used by Shiki transformers
+	 * Raw `meta` information to be used by Shiki transformers.
 	 */
 	meta?: string;
-		},
-	): Promise<string>;
 }
 
 let _cssVariablesTheme: ReturnType<typeof createCssVariablesTheme>;
@@ -31,19 +69,20 @@ export async function createShikiHighlighter({
 	langs = [],
 	theme = 'github-dark',
 	themes = {},
-	defaultColor,
+}: CreateShikiHighlighterOptions = {}): Promise<ShikiHighlighter> {
-	wrap = false,
-	transformers = [],
-}: ShikiConfig = {}): Promise<ShikiHighlighter> {
 	theme = theme === 'css-variables' ? cssVariablesTheme() : theme;
 
-	const highlighter = await getHighlighter({
+	const highlighter = await createHighlighter({
 		langs: ['plaintext', ...langs],
 		themes: Object.values(themes).length ? Object.values(themes) : [theme],
 	});
 
-	return {
+	async function highlight(
-		async highlight(code, lang = 'plaintext', options) {
+		code: string,
+		lang = 'plaintext',
+		options: ShikiHighlighterHighlightOptions,
+		to: 'hast' | 'html',
+	) {
 		const loadedLanguages = highlighter.getLoadedLanguages();
 
 		if (!isSpecialLang(lang) && !loadedLanguages.includes(lang)) {
@@ -51,9 +90,7 @@ export async function createShikiHighlighter({
 				await highlighter.loadLanguage(lang as BundledLanguage);
 			} catch (_err) {
 				// eslint-disable-next-line no-console
-					console.warn(
+				console.warn(`[Shiki] The language "${lang}" doesn't exist, falling back to "plaintext".`);
-						`[Shiki] The language "${lang}" doesn't exist, falling back to "plaintext".`,
-					);
 				lang = 'plaintext';
 			}
 		}
@@ -61,9 +98,9 @@ export async function createShikiHighlighter({
 		const themeOptions = Object.values(themes).length ? { themes } : { theme };
 		const inline = options?.inline ?? false;
 
-			return highlighter.codeToHtml(code, {
+		return highlighter[to === 'html' ? 'codeToHtml' : 'codeToHast'](code, {
 			...themeOptions,
-				defaultColor,
+			defaultColor: options.defaultColor,
 			lang,
 			// NOTE: while we can spread `options.attributes` here so that Shiki can auto-serialize this as rendered
 			// attributes on the top-level tag, it's not clear whether it is fine to pass all attributes as meta, as
@@ -99,9 +136,9 @@ export async function createShikiHighlighter({
 
 						// Handle code wrapping
 						// if wrap=null, do nothing.
-							if (wrap === false) {
+						if (options.wrap === false || options.wrap === undefined) {
 							node.properties.style = styleValue + '; overflow-x: auto;';
-							} else if (wrap === true) {
+						} else if (options.wrap === true) {
 							node.properties.style =
 								styleValue + '; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word;';
 						}
@@ -135,9 +172,17 @@ export async function createShikiHighlighter({
 						}
 					},
 				},
-					...transformers,
+				...(options.transformers ?? []),
 			],
 		});
+	}
+
+	return {
+		codeToHast(code, lang, options = {}) {
+			return highlight(code, lang, options, 'hast') as Promise<Root>;
+		},
+		codeToHtml(code, lang, options = {}) {
+			return highlight(code, lang, options, 'html') as Promise<string>;
 		},
 	};
 }

packages/markdown/remark/src/types.ts

@@ -1,15 +1,10 @@
 import type * as hast from 'hast';
 import type * as mdast from 'mdast';
 import type { Options as RemarkRehypeOptions } from 'remark-rehype';
-import type {
+import type { BuiltinTheme } from 'shiki';
-	BuiltinTheme,
-	LanguageRegistration,
-	ShikiTransformer,
-	ThemeRegistration,
-	ThemeRegistrationRaw,
-} from 'shiki';
 import type * as unified from 'unified';
 import type { DataMap, VFile } from 'vfile';
+import type { CreateShikiHighlighterOptions, ShikiHighlighterHighlightOptions } from './shiki.js';
 
 export type { Node } from 'unist';
 
@@ -35,14 +30,9 @@ export type RemarkRehype = RemarkRehypeOptions;
 
 export type ThemePresets = BuiltinTheme | 'css-variables';
 
-export interface ShikiConfig {
+export interface ShikiConfig
-	langs?: LanguageRegistration[];
+	extends Pick<CreateShikiHighlighterOptions, 'langs' | 'theme' | 'themes'>,
-	theme?: ThemePresets | ThemeRegistration | ThemeRegistrationRaw;
+		Pick<ShikiHighlighterHighlightOptions, 'defaultColor' | 'wrap' | 'transformers'> {}
-	themes?: Record<string, ThemePresets | ThemeRegistration | ThemeRegistrationRaw>;
-	defaultColor?: 'light' | 'dark' | string | false;
-	wrap?: boolean | null;
-	transformers?: ShikiTransformer[];
-}
 
 export interface AstroMarkdownOptions {
 	syntaxHighlight?: 'shiki' | 'prism' | false;

packages/markdown/remark/test/shiki.test.js

@@ -33,16 +33,25 @@ describe('shiki syntax highlighting', () => {
 	it('createShikiHighlighter works', async () => {
 		const highlighter = await createShikiHighlighter();
 
-		const html = await highlighter.highlight('const foo = "bar";', 'js');
+		const html = await highlighter.codeToHtml('const foo = "bar";', 'js');
 
 		assert.match(html, /astro-code github-dark/);
 		assert.match(html, /background-color:#24292e;color:#e1e4e8;/);
 	});
 
+	it('createShikiHighlighter works with codeToHast', async () => {
+		const highlighter = await createShikiHighlighter();
+
+		const hast = await highlighter.codeToHast('const foo = "bar";', 'js');
+
+		assert.match(hast.children[0].properties.class, /astro-code github-dark/);
+		assert.match(hast.children[0].properties.style, /background-color:#24292e;color:#e1e4e8;/);
+	});
+
 	it('diff +/- text has user-select: none', async () => {
 		const highlighter = await createShikiHighlighter();
 
-		const html = await highlighter.highlight(
+		const html = await highlighter.codeToHtml(
 			`\
 - const foo = "bar";
 + const foo = "world";`,
@@ -57,7 +66,7 @@ describe('shiki syntax highlighting', () => {
 	it('renders attributes', async () => {
 		const highlighter = await createShikiHighlighter();
 
-		const html = await highlighter.highlight(`foo`, 'js', {
+		const html = await highlighter.codeToHtml(`foo`, 'js', {
 			attributes: { 'data-foo': 'bar', autofocus: true },
 		});
 
@@ -66,7 +75,10 @@ describe('shiki syntax highlighting', () => {
 	});
 
 	it('supports transformers that reads meta', async () => {
-		const highlighter = await createShikiHighlighter({
+		const highlighter = await createShikiHighlighter();
+
+		const html = await highlighter.codeToHtml(`foo`, 'js', {
+			meta: '{1,3-4}',
 			transformers: [
 				{
 					pre(node) {
@@ -79,10 +91,6 @@ describe('shiki syntax highlighting', () => {
 			],
 		});
 
-		const html = await highlighter.highlight(`foo`, 'js', {
-			meta: '{1,3-4}',
-		});
-
 		assert.match(html, /data-test="\{1,3-4\}"/);
 	});