refactor

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

@@ -148,10 +148,10 @@ declare module '@theme/BlogLayout' {
 }
 
 declare module '@theme/CodeBlock' {
-  import type {ReactElement} from 'react';
+  import type {ReactNode} from 'react';
 
   export interface Props {
-    readonly children: string | ReactElement;
+    readonly children: ReactNode;
     readonly className?: string;
     readonly metastring?: string;
     readonly title?: string;

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

@@ -7,13 +7,12 @@
 
 import React, {
   isValidElement,
-  useEffect,
-  useState,
   type ComponentProps,
-  type ReactElement,
+  type ReactNode,
 } from 'react';
 import clsx from 'clsx';
 import Highlight, {defaultProps, type Language} from 'prism-react-renderer';
+import useIsBrowser from '@docusaurus/useIsBrowser';
 import {
   useThemeConfig,
   parseCodeBlockTitle,
@@ -82,51 +81,47 @@ function CodeBlockLine({
   );
 }
 
-function childrenToString(children: string | ReactElement): string | null {
+/**
+ * Best attempt to make the children a plain string so it is copyable. If there
+ * are react elements, we will not be able to copy the content, and it will
+ * return `children` as-is; otherwise, it concatenates the string children
+ * together.
+ */
+function maybeStringifyChildren(children: ReactNode): ReactNode {
   if (React.Children.toArray(children).some((el) => isValidElement(el))) {
-    return null;
+    return children;
   }
   // The children is now guaranteed to be one/more plain strings
   return Array.isArray(children) ? children.join('') : (children as string);
 }
 
-// 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. There will be a flash seen of the original
-// styles seen using this current approach but that's probably ok. Fixing
-// the flash will require changing the theming approach and is not worth it
-// at this point.
-function useCodeBlockKey(): string {
-  const [mounted, setMounted] = useState(false);
-  useEffect(() => {
-    setMounted(true);
-  }, []);
-  return String(mounted);
-}
-
-export default function CodeBlock({children, ...props}: Props): JSX.Element {
-  const key = useCodeBlockKey();
-  const childrenString = childrenToString(children);
-  return childrenString ? (
-    <CodeBlockString key={key} {...props}>
-      {childrenString}
-    </CodeBlockString>
-  ) : (
-    <CodeBlockJSX key={key} {...props}>
-      {children}
-    </CodeBlockJSX>
+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' ? CodeBlockString : CodeBlockJSX;
+  return (
+    <CodeBlockComp key={String(isBrowser)} {...props}>
+      {children as string}
+    </CodeBlockComp>
   );
 }
 
-function CodeBlockContainer<As extends 'div' | 'pre'>({
-  as: Wrapper,
+function CodeBlockContainer<T extends 'div' | 'pre'>({
+  as: As,
   ...props
-}: {as: As} & ComponentProps<As>) {
+}: {as: T} & ComponentProps<T>) {
   return (
-    // @ts-expect-error: TODO bad typing here
-    <Wrapper
-      {...props}
+    <As
+      // Polymorphic components are hard to type, without `oneOf` generics
+      {...(props as never)}
       className={clsx(
         props.className,
         styles.codeBlockContainer,
@@ -145,7 +140,6 @@ function CodeBlockJSX({children, className}: Props): JSX.Element {
   return (
     <CodeBlockContainer
       as="pre"
-
       tabIndex={0}
       className={clsx(styles.codeBlockStandalone, 'thin-scrollbar', className)}
       style={prismCssVariables}>