portal the registerLanguage api through the lazy component

storybookjs · Jun 22, 2023 · 711f2ea · 711f2ea
1 parent d038135
commit 711f2ea
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/code/ui/components/src/syntaxhighlighter/lazy-syntaxhighlighter.tsx b/code/ui/components/src/syntaxhighlighter/lazy-syntaxhighlighter.tsx
@@ -1,13 +1,29 @@
 import type { ComponentProps } from 'react';
 import React, { Suspense, lazy } from 'react';
 
+import type ReactSyntaxHighlighter from './syntaxhighlighter';
+
+let languages: Parameters<typeof ReactSyntaxHighlighter.registerLanguage>[] = [];
+let Comp: typeof ReactSyntaxHighlighter | null = null;
+
 const LazySyntaxHighlighter = lazy(() => import('./syntaxhighlighter'));
 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>) => (
       <SyntaxHighlighter {...props} formatter={formatter} />
@@ -24,3 +40,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);
+};