Codeblock cleanup in markdown

reflex/components/datadisplay/code.py

@@ -10,6 +10,7 @@ from reflex.components.core.cond import color_mode_cond
 from reflex.components.lucide.icon import Icon
 from reflex.components.radix.themes.components.button import Button
 from reflex.components.radix.themes.layout.box import Box
+from reflex.components.markdown.markdown import MarkdownComponentMapMixin, _PROPS, _CHILDREN
 from reflex.constants.colors import Color
 from reflex.event import set_clipboard
 from reflex.style import Style
@@ -378,7 +379,7 @@ for theme_name in dir(Theme):
     setattr(Theme, theme_name, getattr(Theme, theme_name)._replace(_var_type=Theme))
 
 
-class CodeBlock(Component):
+class CodeBlock(Component, MarkdownComponentMapMixin):
     """A code block."""
 
     library = "react-syntax-highlighter@15.6.1"
@@ -543,6 +544,28 @@ class CodeBlock(Component):
     def _exclude_props(self) -> list[str]:
         return ["can_copy", "copy_button"]
 
+    @classmethod
+    def get_component_map_custom_code(cls) -> str:
+        """Get the custom code for the component.
+
+        Returns:
+            The custom code for the component.
+        """
+        return """
+const match = (className || '').match(/language-(?<lang>.*)/);
+    const language = match ? match[1] : '';
+    if (language) {
+    (async () => {
+      try {
+        const module = await import(`react-syntax-highlighter/dist/cjs/languages/prism/${{language}}`);
+        SyntaxHighlighter.registerLanguage(language, module.default);
+      } catch (error) {
+        console.error(`Error importing language module for ${language}:`, error);
+      }
+    })();
+  }       
+        """
+
 
 class CodeblockNamespace(ComponentNamespace):
     """Namespace for the CodeBlock component."""

reflex/components/markdown/markdown.py

@@ -74,6 +74,11 @@ def get_base_component_map() -> dict[str, Callable]:
     }
 
 
+class MarkdownComponentMapMixin:
+    def get_component_map_custom_code(self) -> str:
+        return ""
+
+
 class Markdown(Component):
     """A markdown component."""
 
@@ -179,8 +184,8 @@ class Markdown(Component):
                 component(_MOCK_ARG)._get_all_imports()  # type: ignore
                 for component in self.component_map.values()
             ],
-            CodeBlock.create(theme=Theme.light)._get_imports(),
+            # CodeBlock.create(theme=Theme.light)._get_imports(),
-            Code.create()._get_imports(),
+            # Code.create()._get_imports(),
         ]
 
     def get_component(self, tag: str, **props) -> Component:
@@ -251,22 +256,13 @@ class Markdown(Component):
             )
             for tag in self.component_map
         }
-
+        codeblock_component = self.get_component("codeblock")
+        codeblock_custom_code = codeblock_component.get_component_map_custom_code() if hasattr(codeblock_component,
+                                                                                               "get_component_map_custom_code") else ""
         # Separate out inline code and code blocks.
         components["code"] = Var(
             _js_expr=f"""(({{node, inline, className, {_CHILDREN._js_expr}, {_PROPS._js_expr}}}) => {{
-    const match = (className || '').match(/language-(?<lang>.*)/);
+    {codeblock_custom_code};
-    const language = match ? match[1] : '';
-    if (language) {{
-    (async () => {{
-      try {{
-        const module = await import(`react-syntax-highlighter/dist/cjs/languages/prism/${{language}}`);
-        SyntaxHighlighter.registerLanguage(language, module.default);
-      }} catch (error) {{
-        console.error(`Error importing language module for ${{language}}:`, error);
-      }}
-    }})();
-  }}
     return inline ? (
         {self.format_component("code")}
     ) : (