Lazy loaded shiki languages during syntax highlighting (#10618)

43081j · web-flow · commit 374efcdff962 · 2024-04-01T22:52:50.000+08:00
diff --git a/.changeset/real-rabbits-bake.md b/.changeset/real-rabbits-bake.md
@@ -0,0 +1,5 @@
+---
+"@astrojs/markdown-remark": major
+---
+
+Updates Shiki syntax highlighting to lazily load shiki languages by default (only preloading `plaintext`). Additionally, the `createShikiHighlighter()` API now returns an asynchronous `highlight()` function due to this.
diff --git a/packages/astro/components/Code.astro b/packages/astro/components/Code.astro
@@ -87,7 +87,7 @@ const highlighter = await getCachedHighlighter({
 	wrap,
 });
 
-const html = highlighter.highlight(code, typeof lang === 'string' ? lang : lang.name, {
+const html = await highlighter.highlight(code, typeof lang === 'string' ? lang : lang.name, {
 	inline,
 	attributes: rest as any,
 });
diff --git a/packages/astro/src/assets/vite-plugin-assets.ts b/packages/astro/src/assets/vite-plugin-assets.ts
@@ -2,7 +2,7 @@ import { extname } from 'node:path';
 import MagicString from 'magic-string';
 import type * as vite from 'vite';
 import { normalizePath } from 'vite';
-import type { AstroPluginOptions, ImageTransform } from '../@types/astro.js';
+import type { AstroPluginOptions, AstroSettings, ImageTransform } from '../@types/astro.js';
 import { extendManualChunks } from '../core/build/plugins/util.js';
 import { AstroError, AstroErrorData } from '../core/errors/index.js';
 import {
@@ -24,6 +24,71 @@ const resolvedVirtualModuleId = '\0' + VIRTUAL_MODULE_ID;
 
 const assetRegex = new RegExp(`\\.(${VALID_INPUT_FORMATS.join('|')})`, 'i');
 const assetRegexEnds = new RegExp(`\\.(${VALID_INPUT_FORMATS.join('|')})$`, 'i');
+const addStaticImageFactory = (
+	settings: AstroSettings
+): typeof globalThis.astroAsset.addStaticImage => {
+	return (options, hashProperties, originalFSPath) => {
+		if (!globalThis.astroAsset.staticImages) {
+			globalThis.astroAsset.staticImages = new Map<
+				string,
+				{
+					originalSrcPath: string;
+					transforms: Map<string, { finalPath: string; transform: ImageTransform }>;
+				}
+			>();
+		}
+
+		// Rollup will copy the file to the output directory, as such this is the path in the output directory, including the asset prefix / base
+		const ESMImportedImageSrc = isESMImportedImage(options.src) ? options.src.src : options.src;
+		const fileExtension = extname(ESMImportedImageSrc);
+		const assetPrefix = getAssetsPrefix(fileExtension, settings.config.build.assetsPrefix);
+
+		// This is the path to the original image, from the dist root, without the base or the asset prefix (e.g. /_astro/image.hash.png)
+		const finalOriginalPath = removeBase(
+			removeBase(ESMImportedImageSrc, settings.config.base),
+			assetPrefix
+		);
+
+		const hash = hashTransform(options, settings.config.image.service.entrypoint, hashProperties);
+
+		let finalFilePath: string;
+		let transformsForPath = globalThis.astroAsset.staticImages.get(finalOriginalPath);
+		let transformForHash = transformsForPath?.transforms.get(hash);
+
+		// If the same image has already been transformed with the same options, we'll reuse the final path
+		if (transformsForPath && transformForHash) {
+			finalFilePath = transformForHash.finalPath;
+		} else {
+			finalFilePath = prependForwardSlash(
+				joinPaths(
+					isESMImportedImage(options.src) ? '' : settings.config.build.assets,
+					prependForwardSlash(propsToFilename(finalOriginalPath, options, hash))
+				)
+			);
+
+			if (!transformsForPath) {
+				globalThis.astroAsset.staticImages.set(finalOriginalPath, {
+					originalSrcPath: originalFSPath,
+					transforms: new Map(),
+				});
+				transformsForPath = globalThis.astroAsset.staticImages.get(finalOriginalPath)!;
+			}
+
+			transformsForPath.transforms.set(hash, {
+				finalPath: finalFilePath,
+				transform: options,
+			});
+		}
+
+		// The paths here are used for URLs, so we need to make sure they have the proper format for an URL
+		// (leading slash, prefixed with the base / assets prefix, encoded, etc)
+		if (settings.config.build.assetsPrefix) {
+			return encodeURI(joinPaths(assetPrefix, finalFilePath));
+		} else {
+			return encodeURI(prependForwardSlash(joinPaths(settings.config.base, finalFilePath)));
+		}
+	};
+};
 
 export default function assets({
 	settings,
@@ -92,73 +157,7 @@ export default function assets({
 					return;
 				}
 
-				globalThis.astroAsset.addStaticImage = (options, hashProperties, originalFSPath) => {
-					if (!globalThis.astroAsset.staticImages) {
-						globalThis.astroAsset.staticImages = new Map<
-							string,
-							{
-								originalSrcPath: string;
-								transforms: Map<string, { finalPath: string; transform: ImageTransform }>;
-							}
-						>();
-					}
-
-					// Rollup will copy the file to the output directory, as such this is the path in the output directory, including the asset prefix / base
-					const ESMImportedImageSrc = isESMImportedImage(options.src)
-						? options.src.src
-						: options.src;
-					const fileExtension = extname(ESMImportedImageSrc);
-					const assetPrefix = getAssetsPrefix(fileExtension, settings.config.build.assetsPrefix);
-
-					// This is the path to the original image, from the dist root, without the base or the asset prefix (e.g. /_astro/image.hash.png)
-					const finalOriginalPath = removeBase(
-						removeBase(ESMImportedImageSrc, settings.config.base),
-						assetPrefix
-					);
-
-					const hash = hashTransform(
-						options,
-						settings.config.image.service.entrypoint,
-						hashProperties
-					);
-
-					let finalFilePath: string;
-					let transformsForPath = globalThis.astroAsset.staticImages.get(finalOriginalPath);
-					let transformForHash = transformsForPath?.transforms.get(hash);
-
-					// If the same image has already been transformed with the same options, we'll reuse the final path
-					if (transformsForPath && transformForHash) {
-						finalFilePath = transformForHash.finalPath;
-					} else {
-						finalFilePath = prependForwardSlash(
-							joinPaths(
-								isESMImportedImage(options.src) ? '' : settings.config.build.assets,
-								prependForwardSlash(propsToFilename(finalOriginalPath, options, hash))
-							)
-						);
-
-						if (!transformsForPath) {
-							globalThis.astroAsset.staticImages.set(finalOriginalPath, {
-								originalSrcPath: originalFSPath,
-								transforms: new Map(),
-							});
-							transformsForPath = globalThis.astroAsset.staticImages.get(finalOriginalPath)!;
-						}
-
-						transformsForPath.transforms.set(hash, {
-							finalPath: finalFilePath,
-							transform: options,
-						});
-					}
-
-					// The paths here are used for URLs, so we need to make sure they have the proper format for an URL
-					// (leading slash, prefixed with the base / assets prefix, encoded, etc)
-					if (settings.config.build.assetsPrefix) {
-						return encodeURI(joinPaths(assetPrefix, finalFilePath));
-					} else {
-						return encodeURI(prependForwardSlash(joinPaths(settings.config.base, finalFilePath)));
-					}
-				};
+				globalThis.astroAsset.addStaticImage = addStaticImageFactory(settings);
 			},
 			// In build, rewrite paths to ESM imported images in code to their final location
 			async renderChunk(code) {
diff --git a/packages/astro/src/content/vite-plugin-content-assets.ts b/packages/astro/src/content/vite-plugin-content-assets.ts
@@ -257,6 +257,8 @@ export function astroConfigBuildPlugin(
 						mutate(chunk, ['server'], newCode);
 					}
 				}
+
+				ssrPluginContext = undefined;
 			},
 		},
 	};
diff --git a/packages/astro/src/vite-plugin-markdown/index.ts b/packages/astro/src/vite-plugin-markdown/index.ts
@@ -32,14 +32,17 @@ const astroErrorModulePath = normalizePath(
 );
 
 export default function markdown({ settings, logger }: AstroPluginOptions): Plugin {
-	let processor: MarkdownProcessor;
+	let processor: MarkdownProcessor | undefined;
 
 	return {
 		enforce: 'pre',
 		name: 'astro:markdown',
 		async buildStart() {
 			processor = await createMarkdownProcessor(settings.config.markdown);
 		},
+		buildEnd() {
+			processor = undefined;
+		},
 		// Why not the "transform" hook instead of "load" + readFile?
 		// A: Vite transforms all "import.meta.env" references to their values before
 		// passing to the transform hook. This lets us get the truly raw value
@@ -52,7 +55,7 @@ export default function markdown({ settings, logger }: AstroPluginOptions): Plug
 
 				const fileURL = pathToFileURL(fileId);
 
-				const renderResult = await processor
+				const renderResult = await processor!
 					.render(raw.content, {
 						// @ts-expect-error passing internal prop
 						fileURL,
diff --git a/packages/astro/test/astro-markdown-shiki.test.js b/packages/astro/test/astro-markdown-shiki.test.js
@@ -80,26 +80,42 @@ describe('Astro Markdown Shiki', () => {
 		});
 	});
 
-	describe('Custom langs', () => {
+	describe('Languages', () => {
 		let fixture;
+		let $;
 
 		before(async () => {
 			fixture = await loadFixture({ root: './fixtures/astro-markdown-shiki/langs/' });
 			await fixture.build();
-		});
-
-		it('Markdown file', async () => {
 			const html = await fixture.readFile('/index.html');
-			const $ = cheerio.load(html);
+			$ = cheerio.load(html);
+		});
 
-			const segments = $('.line').get(6).children;
+		it('custom language', async () => {
+			const lang = $('.astro-code').get(0);
+			const segments = $('.line', lang).get(6).children;
 			assert.equal(segments.length, 2);
 			assert.equal(segments[0].attribs.style, 'color:#79B8FF');
 			assert.equal(segments[1].attribs.style, 'color:#E1E4E8');
+		});
 
+		it('handles unknown languages', () => {
 			const unknownLang = $('.astro-code').get(1);
 			assert.ok(unknownLang.attribs.style.includes('background-color:#24292e;color:#e1e4e8;'));
 		});
+
+		it('handles lazy loaded languages', () => {
+			const lang = $('.astro-code').get(2);
+			const segments = $('.line', lang).get(0).children;
+			assert.equal(segments.length, 7);
+			assert.equal(segments[0].attribs.style, 'color:#F97583');
+			assert.equal(segments[1].attribs.style, 'color:#79B8FF');
+			assert.equal(segments[2].attribs.style, 'color:#F97583');
+			assert.equal(segments[3].attribs.style, 'color:#79B8FF');
+			assert.equal(segments[4].attribs.style, 'color:#F97583');
+			assert.equal(segments[5].attribs.style, 'color:#79B8FF');
+			assert.equal(segments[6].attribs.style, 'color:#E1E4E8');
+		});
 	});
 
 	describe('Wrapping behaviours', () => {
diff --git a/packages/astro/test/fixtures/astro-markdown-shiki/langs/src/pages/index.md b/packages/astro/test/fixtures/astro-markdown-shiki/langs/src/pages/index.md
@@ -24,3 +24,7 @@ fin
 ```unknown
 This language does not exist
 ```
+
+```ts
+const someTypeScript: number = 5;
+```
diff --git a/packages/create-astro/test/fixtures/not-empty/package.json b/packages/create-astro/test/fixtures/not-empty/package.json
@@ -6,4 +6,4 @@
     "build": "astro build",
     "preview": "astro preview"
   }
-}
+}
diff --git a/packages/integrations/markdoc/components/Renderer.astro b/packages/integrations/markdoc/components/Renderer.astro
@@ -12,7 +12,7 @@ type Props = {
 const { stringifiedAst, config } = Astro.props as Props;
 
 const ast = Markdoc.Ast.fromJSON(stringifiedAst);
-const content = Markdoc.transform(ast, config);
+const content = await Markdoc.transform(ast, config);
 ---
 
 {
diff --git a/packages/integrations/markdoc/src/extensions/shiki.ts b/packages/integrations/markdoc/src/extensions/shiki.ts
@@ -11,12 +11,12 @@ export default async function shiki(config?: ShikiConfig): Promise<AstroMarkdocC
 		nodes: {
 			fence: {
 				attributes: Markdoc.nodes.fence.attributes!,
-				transform({ attributes }) {
+				async transform({ attributes }) {
 					// NOTE: The `meta` from fence code, e.g. ```js {1,3-4}, isn't quite supported by Markdoc.
 					// 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 = highlighter.highlight(attributes.content, lang);
+					const html = await highlighter.highlight(attributes.content, lang);
 
 					// Use `unescapeHTML` to return `HTMLString` for Astro renderer to inline as HTML
 					return unescapeHTML(html) as any;
diff --git a/packages/integrations/markdoc/test/syntax-highlighting.test.js b/packages/integrations/markdoc/test/syntax-highlighting.test.js
@@ -23,7 +23,7 @@ describe('Markdoc - syntax highlighting', () => {
 	describe('shiki', () => {
 		it('transforms with defaults', async () => {
 			const ast = Markdoc.parse(entry);
-			const content = Markdoc.transform(ast, await getConfigExtendingShiki());
+			const content = await Markdoc.transform(ast, await getConfigExtendingShiki());
 
 			assert.equal(content.children.length, 2);
 			for (const codeBlock of content.children) {
@@ -36,7 +36,7 @@ describe('Markdoc - syntax highlighting', () => {
 		});
 		it('transforms with `theme` property', async () => {
 			const ast = Markdoc.parse(entry);
-			const content = Markdoc.transform(
+			const content = await Markdoc.transform(
 				ast,
 				await getConfigExtendingShiki({
 					theme: 'dracula',
@@ -53,7 +53,7 @@ describe('Markdoc - syntax highlighting', () => {
 		});
 		it('transforms with `wrap` property', async () => {
 			const ast = Markdoc.parse(entry);
-			const content = Markdoc.transform(
+			const content = await Markdoc.transform(
 				ast,
 				await getConfigExtendingShiki({
 					wrap: true,
@@ -76,7 +76,7 @@ describe('Markdoc - syntax highlighting', () => {
 			const config = await setupConfig({
 				extends: [prism()],
 			});
-			const content = Markdoc.transform(ast, config);
+			const content = await Markdoc.transform(ast, config);
 
 			assert.equal(content.children.length, 2);
 			const [tsBlock, cssBlock] = content.children;
diff --git a/packages/markdown/remark/src/highlight.ts b/packages/markdown/remark/src/highlight.ts
@@ -1,10 +1,10 @@
-import type { Element, Root } from 'hast';
+import type { Element, Parent, Root } from 'hast';
 import { fromHtml } from 'hast-util-from-html';
 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 }) => string;
+type Highlighter = (code: string, language: string, options?: { meta?: string }) => Promise<string>;
 
 const languagePattern = /\blanguage-(\S+)\b/;
 
@@ -17,7 +17,14 @@ const languagePattern = /\blanguage-(\S+)\b/;
  *   A fnction which receives the code and language, and returns the HTML of a syntax
  *   highlighted `<pre>` element.
  */
-export function highlightCodeBlocks(tree: Root, highlighter: Highlighter) {
+export async function highlightCodeBlocks(tree: Root, highlighter: Highlighter) {
+	const nodes: Array<{
+		node: Element;
+		language: string;
+		parent: Element;
+		grandParent: Parent;
+	}> = [];
+
 	// We’re looking for `<code>` elements
 	visitParents(tree, { type: 'element', tagName: 'code' }, (node, ancestors) => {
 		const parent = ancestors.at(-1);
@@ -55,17 +62,25 @@ export function highlightCodeBlocks(tree: Root, highlighter: Highlighter) {
 			return;
 		}
 
+		nodes.push({
+			node,
+			language: languageMatch?.[1] || 'plaintext',
+			parent,
+			grandParent: ancestors.at(-2)!,
+		});
+	});
+
+	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' });
-		const html = highlighter(code, languageMatch?.[1] || 'plaintext', { meta });
+		const html = await highlighter(code, language, { meta });
 		// The replacement returns a root node with 1 child, the `<pr>` element replacement.
 		const replacement = fromHtml(html, { fragment: true }).children[0] as Element;
 		// We just generated this node, so any positional information is invalid.
 		removePosition(replacement);
 
 		// We replace the parent in its parent with the new `<pre>` element.
-		const grandParent = ancestors.at(-2)!;
 		const index = grandParent.children.indexOf(parent);
 		grandParent.children[index] = replacement;
-	});
+	}
 }
diff --git a/packages/markdown/remark/src/rehype-prism.ts b/packages/markdown/remark/src/rehype-prism.ts
diff --git a/packages/markdown/remark/src/rehype-shiki.ts b/packages/markdown/remark/src/rehype-shiki.ts
diff --git a/packages/markdown/remark/src/shiki.ts b/packages/markdown/remark/src/shiki.ts
diff --git a/packages/markdown/remark/test/shiki.test.js b/packages/markdown/remark/test/shiki.test.js

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@astrojs/markdown-remark": major
 +---
++
 +Updates Shiki syntax highlighting to lazily load shiki languages by default (only preloading `plaintext`). Additionally, the `createShikiHighlighter()` API now returns an asynchronous `highlight()` function due to this.
Original file line number	Diff line number	Diff line change
`@@ -257,6 +257,8 @@ export function astroConfigBuildPlugin(`
`257`	`257`	`mutate(chunk, ['server'], newCode);`
`258`	`258`	`}`
`259`	`259`	`}`
	`260`	`+`
	`261`	`+ ssrPluginContext = undefined;`
`260`	`262`	`},`
`261`	`263`	`},`
`262`	`264`	`};`
Original file line number	Diff line number	Diff line change
`@@ -6,4 +6,4 @@`
`6`	`6`	`"build": "astro build",`
`7`	`7`	`"preview": "astro preview"`
`8`	`8`	`}`
`9`		`-}`
	`9`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ type Props = {`
`12`	`12`	`const { stringifiedAst, config } = Astro.props as Props;`
`13`	`13`
`14`	`14`	`const ast = Markdoc.Ast.fromJSON(stringifiedAst);`
`15`		`-const content = Markdoc.transform(ast, config);`
	`15`	`+const content = await Markdoc.transform(ast, config);`
`16`	`16`	`---`
`17`	`17`
`18`	`18`	`{`