SyntaxHighlighter: Expose registerLanguage

code/ui/blocks/src/components/Source.tsx

@@ -1,27 +1,26 @@
 import type { ComponentProps, FunctionComponent } from 'react';
+import { ThemeProvider, convert, ignoreSsrWarning, styled, themes } from '@storybook/theming';
+
 import React from 'react';
-import { styled, ThemeProvider, convert, themes, ignoreSsrWarning } from '@storybook/theming';
 import { SyntaxHighlighter } from '@storybook/components';
-
+import type { SyntaxHighlighterProps } from '@storybook/components';
 import { EmptyBlock } from './EmptyBlock';
 
-const StyledSyntaxHighlighter: typeof SyntaxHighlighter = styled(SyntaxHighlighter)(
-  ({ theme }) => ({
-    // DocBlocks-specific styling and overrides
-    fontSize: `${theme.typography.size.s2 - 1}px`,
-    lineHeight: '19px',
-    margin: '25px 0 40px',
-    borderRadius: theme.appBorderRadius,
-    boxShadow:
-      theme.base === 'light'
-        ? 'rgba(0, 0, 0, 0.10) 0 1px 3px 0'
-        : 'rgba(0, 0, 0, 0.20) 0 2px 5px 0',
-    'pre.prismjs': {
-      padding: 20,
-      background: 'inherit',
-    },
-  })
-);
+const StyledSyntaxHighlighter: React.FunctionComponent<SyntaxHighlighterProps> = styled(
+  SyntaxHighlighter
+)(({ theme }) => ({
+  // DocBlocks-specific styling and overrides
+  fontSize: `${theme.typography.size.s2 - 1}px`,
+  lineHeight: '19px',
+  margin: '25px 0 40px',
+  borderRadius: theme.appBorderRadius,
+  boxShadow:
+    theme.base === 'light' ? 'rgba(0, 0, 0, 0.10) 0 1px 3px 0' : 'rgba(0, 0, 0, 0.20) 0 2px 5px 0',
+  'pre.prismjs': {
+    padding: 20,
+    background: 'inherit',
+  },
+}));
 
 export enum SourceError {
   NO_STORY = 'There\u2019s no story here.',

code/ui/components/src/syntaxhighlighter/lazy-syntaxhighlighter.tsx

@@ -1,21 +1,57 @@
-import type { ComponentProps } from 'react';
 import React, { Suspense, lazy } from 'react';
 
-const LazySyntaxHighlighter = lazy(() => import('./syntaxhighlighter'));
+import type { ComponentProps } from 'react';
+import type ReactSyntaxHighlighter from './syntaxhighlighter';
+
+let languages: Parameters<typeof ReactSyntaxHighlighter.registerLanguage>[] = [];
+let Comp: typeof ReactSyntaxHighlighter | null = null;
+
+const LazySyntaxHighlighter = lazy(async () => {
+  const { SyntaxHighlighter } = await import('./syntaxhighlighter');
+
+  if (languages.length > 0) {
+    languages.forEach((args) => {
+      SyntaxHighlighter.registerLanguage(...args);
+    });
+    languages = [];
+  }
+
+  if (Comp === null) Comp = SyntaxHighlighter;
+
+  return {
+    default: (props: ComponentProps<typeof SyntaxHighlighter>) => <SyntaxHighlighter {...props} />,
+  };
+});
+
 const LazySyntaxHighlighterWithFormatter = lazy(async () => {
   const [{ SyntaxHighlighter }, { formatter }] = await Promise.all([
     import('./syntaxhighlighter'),
     import('./formatter'),
   ]);
 
+  if (languages.length > 0) {
+    languages.forEach((args) => {
+      SyntaxHighlighter.registerLanguage(...args);
+    });
+    languages = [];
+  }
+
+  if (Comp === null) {
+    Comp = SyntaxHighlighter;
+  }
+
   return {
-    default: (props: ComponentProps<typeof LazySyntaxHighlighter>) => (
+    default: (props: ComponentProps<typeof SyntaxHighlighter>) => (
       <SyntaxHighlighter {...props} formatter={formatter} />
     ),
   };
 });
 
-export const SyntaxHighlighter = (props: ComponentProps<typeof LazySyntaxHighlighter>) => (
+export const SyntaxHighlighter = (
+  props:
+    | ComponentProps<typeof LazySyntaxHighlighter>
+    | ComponentProps<typeof LazySyntaxHighlighterWithFormatter>
+) => (
   <Suspense fallback={<div />}>
     {props.format !== false ? (
       <LazySyntaxHighlighterWithFormatter {...props} />
@@ -24,3 +60,13 @@ export const SyntaxHighlighter = (props: ComponentProps<typeof LazySyntaxHighlig
     )}
   </Suspense>
 );
+
+SyntaxHighlighter.registerLanguage = (
+  ...args: Parameters<typeof ReactSyntaxHighlighter.registerLanguage>
+) => {
+  if (Comp !== null) {
+    Comp.registerLanguage(...args);
+    return;
+  }
+  languages.push(args);
+};

code/ui/components/src/syntaxhighlighter/syntaxhighlighter.stories.tsx

@@ -1,6 +1,7 @@
+import { ThemeProvider, ensure, themes } from '@storybook/theming';
+
 import type { ComponentProps } from 'react';
 import React from 'react';
-import { ThemeProvider, themes, ensure } from '@storybook/theming';
 import { SyntaxHighlighter } from './lazy-syntaxhighlighter';
 
 export default {
@@ -108,6 +109,24 @@ export const GraphQL = {
   },
 };
 
+export const CustomSyntax = {
+  args: {
+    language: 'scss',
+    children: `// Custom language syntax registered
+div.parent {
+  div.child {
+    color: $red;
+  }
+}`,
+  },
+  loaders: [
+    async () => {
+      const scss = (await import('react-syntax-highlighter/dist/esm/languages/prism/scss')).default;
+      SyntaxHighlighter.registerLanguage('scss', scss);
+    },
+  ],
+};
+
 export const Unsupported = {
   args: {
     language: 'C#',

code/ui/components/src/syntaxhighlighter/syntaxhighlighter.tsx

@@ -189,7 +189,7 @@ export interface SyntaxHighlighterState {
 
 // copied from @types/react-syntax-highlighter/index.d.ts
 
-export const SyntaxHighlighter: FC<SyntaxHighlighterProps> = ({
+export const SyntaxHighlighter = ({
   children,
   language = 'jsx',
   copyable = false,
@@ -200,7 +200,7 @@ export const SyntaxHighlighter: FC<SyntaxHighlighterProps> = ({
   className = null,
   showLineNumbers = false,
   ...rest
-}) => {
+}: SyntaxHighlighterProps) => {
   if (typeof children !== 'string' || !children.trim()) {
     return null;
   }
@@ -254,4 +254,8 @@ export const SyntaxHighlighter: FC<SyntaxHighlighterProps> = ({
   );
 };
 
+SyntaxHighlighter.registerLanguage = (
+  ...args: Parameters<typeof ReactSyntaxHighlighter.registerLanguage>
+) => ReactSyntaxHighlighter.registerLanguage(...args);
+
 export default SyntaxHighlighter;