refactor: mitigate FOUC when applying Prism theme

packages/docusaurus-theme-classic/src/index.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import type {CSSProperties} from 'react';
 import path from 'path';
 import {createRequire} from 'module';
 import rtlcss from 'rtlcss';
@@ -15,13 +16,30 @@ import type {ThemeConfig} from '@docusaurus/theme-common';
 import type {Plugin as PostCssPlugin} from 'postcss';
 import type {Options} from '@docusaurus/theme-classic';
 import type webpack from 'webpack';
+import type {PrismTheme} from 'prism-react-renderer';
 
 const requireFromDocusaurusCore = createRequire(
   require.resolve('@docusaurus/core/package.json'),
 );
 const ContextReplacementPlugin: typeof webpack.ContextReplacementPlugin =
   requireFromDocusaurusCore('webpack/lib/ContextReplacementPlugin');
 
+const getPrismCssVariables = (prismTheme: PrismTheme): CSSProperties => {
+  const mapping: {[name: keyof PrismTheme['plain']]: string} = {
+    color: '--prism-color',
+    backgroundColor: '--prism-background-color',
+  };
+
+  const properties: {[key: string]: string} = {};
+  Object.entries(prismTheme.plain).forEach(([key, value]) => {
+    const varName = mapping[key];
+    if (varName && typeof value === 'string') {
+      properties[varName] = value;
+    }
+  });
+  return properties;
+};
+
 // Need to be inlined to prevent dark mode FOUC
 // Make sure the key is the same as the one in `/theme/hooks/useTheme.js`
 const ThemeStorageKey = 'theme';
@@ -103,10 +121,14 @@ export default function themeClassic(
   const {
     announcementBar,
     colorMode,
-    prism: {additionalLanguages},
+    prism: {additionalLanguages, theme, darkTheme},
   } = themeConfig;
   const {customCss} = options ?? {};
   const {direction} = localeConfigs[currentLocale]!;
+  const prismBaseStyles = {
+    ':root': getPrismCssVariables(theme),
+    '[data-theme="dark"]': getPrismCssVariables(darkTheme),
+  };
 
   return {
     name: 'docusaurus-theme-classic',
@@ -201,6 +223,17 @@ ${noFlashColorMode(colorMode)}
 ${announcementBar ? AnnouncementBarInlineJavaScript : ''}
             `,
           },
+          {
+            tagName: 'style',
+            innerHTML: Object.entries(prismBaseStyles)
+              .map(
+                ([selector, properties]) =>
+                  `${selector} {${Object.entries(properties)
+                    .map(([name, value]) => `${name}:${value};`)
+                    .join('')}}`,
+              )
+              .join(' '),
+          },
         ],
       };
     },

packages/docusaurus-theme-classic/src/theme/CodeBlock/Container/index.tsx

@@ -7,24 +7,17 @@
 
 import React, {type ComponentProps} from 'react';
 import clsx from 'clsx';
-import {
-  usePrismTheme,
-  getPrismCssVariables,
-  ThemeClassNames,
-} from '@docusaurus/theme-common';
+import {ThemeClassNames} from '@docusaurus/theme-common';
 import styles from './styles.module.css';
 
 export default function CodeBlockContainer<T extends 'div' | 'pre'>({
   as: As,
   ...props
 }: {as: T} & ComponentProps<T>): JSX.Element {
-  const prismTheme = usePrismTheme();
-  const prismCssVariables = getPrismCssVariables(prismTheme);
   return (
     <As
       // Polymorphic components are hard to type, without `oneOf` generics
       {...(props as never)}
-      style={prismCssVariables}
       className={clsx(
         props.className,
         styles.codeBlockContainer,

packages/docusaurus-theme-classic/src/theme/CodeBlock/Line/index.tsx

@@ -7,6 +7,7 @@
 
 import React from 'react';
 import clsx from 'clsx';
+import useIsBrowser from '@docusaurus/useIsBrowser';
 import type {Props} from '@theme/CodeBlock/Line';
 
 import styles from './styles.module.css';
@@ -18,6 +19,8 @@ export default function CodeBlockLine({
   getLineProps,
   getTokenProps,
 }: Props): JSX.Element {
+  const isBrowser = useIsBrowser();
+
   if (line.length === 1 && line[0]!.content === '\n') {
     line[0]!.content = '';
   }
@@ -28,11 +31,15 @@ export default function CodeBlockLine({
   });
 
   const lineTokens = line.map((token, key) => (
-    <span key={key} {...getTokenProps({token, key})} />
+    <span
+      key={key}
+      {...getTokenProps({token, key})}
+      {...(!isBrowser && {style: undefined})}
+    />
   ));
 
   return (
-    <span {...lineProps}>
+    <span {...lineProps} {...(!isBrowser && {style: undefined})}>
       {showLineNumbers ? (
         <>
           <span className={styles.codeLineNumber} />

packages/docusaurus-theme-classic/src/theme/CodeBlock/index.tsx

@@ -6,7 +6,6 @@
  */
 
 import React, {isValidElement, type ReactNode} from 'react';
-import useIsBrowser from '@docusaurus/useIsBrowser';
 import ElementContent from '@theme/CodeBlock/Content/Element';
 import StringContent from '@theme/CodeBlock/Content/String';
 import type {Props} from '@theme/CodeBlock';
@@ -29,17 +28,8 @@ export default function CodeBlock({
   children: rawChildren,
   ...props
 }: Props): JSX.Element {
-  // The Prism theme on SSR is always the default theme but the site theme can
-  // be in a different mode. React hydration doesn't update DOM styles that come
-  // from SSR. Hence force a re-render after mounting to apply the current
-  // relevant styles.
-  const isBrowser = useIsBrowser();
   const children = maybeStringifyChildren(rawChildren);
   const CodeBlockComp =
     typeof children === 'string' ? StringContent : ElementContent;
-  return (
-    <CodeBlockComp key={String(isBrowser)} {...props}>
-      {children as string}
-    </CodeBlockComp>
-  );
+  return <CodeBlockComp {...props}>{children as string}</CodeBlockComp>;
 }

packages/docusaurus-theme-classic/src/validateThemeConfig.ts

@@ -44,6 +44,7 @@ export const DEFAULT_CONFIG = {
   prism: {
     additionalLanguages: [],
     theme: defaultPrismTheme,
+    darkTheme: defaultPrismTheme,
     magicComments: [
       {
         className: 'theme-code-block-highlighted-line',

packages/docusaurus-theme-common/src/hooks/usePrismTheme.ts

@@ -16,9 +16,7 @@ import type {PrismTheme} from 'prism-react-renderer';
 export function usePrismTheme(): PrismTheme {
   const {prism} = useThemeConfig();
   const {colorMode} = useColorMode();
-  const lightModeTheme = prism.theme;
-  const darkModeTheme = prism.darkTheme || lightModeTheme;
-  const prismTheme = colorMode === 'dark' ? darkModeTheme : lightModeTheme;
+  const prismTheme = colorMode === 'dark' ? prism.darkTheme : prism.theme;
 
   return prismTheme;
 }

packages/docusaurus-theme-common/src/index.ts

@@ -35,7 +35,6 @@ export {
   parseLanguage,
   parseLines,
   containsLineNumbers,
-  getPrismCssVariables,
 } from './utils/codeBlockUtils';
 
 export {

packages/docusaurus-theme-common/src/utils/codeBlockUtils.ts

@@ -5,9 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import type {CSSProperties} from 'react';
 import rangeParser from 'parse-numeric-range';
-import type {PrismTheme} from 'prism-react-renderer';
 
 const codeBlockTitleRegex = /title=(?<quote>["'])(?<title>.*?)\1/;
 const metastringLinesRangeRegex = /\{(?<range>[\d,-]+)\}/;
@@ -231,19 +229,3 @@ export function parseLines(
   });
   return {lineClassNames, code};
 }
-
-export function getPrismCssVariables(prismTheme: PrismTheme): CSSProperties {
-  const mapping: {[name: keyof PrismTheme['plain']]: string} = {
-    color: '--prism-color',
-    backgroundColor: '--prism-background-color',
-  };
-
-  const properties: {[key: string]: string} = {};
-  Object.entries(prismTheme.plain).forEach(([key, value]) => {
-    const varName = mapping[key];
-    if (varName && typeof value === 'string') {
-      properties[varName] = value;
-    }
-  });
-  return properties;
-}

packages/docusaurus-theme-common/src/utils/useThemeConfig.ts

@@ -55,7 +55,7 @@ export type AnnouncementBarConfig = {
 
 export type PrismConfig = {
   theme: PrismTheme;
-  darkTheme?: PrismTheme;
+  darkTheme: PrismTheme;
   defaultLanguage?: string;
   additionalLanguages: string[];
   magicComments: MagicCommentConfig[];