[ENG-4010]Codeblock cleanup in markdown (#4233)

* Codeblock cleanup in markdown

* Initial approach to getting this working with rx.memo and reflex web

* abstract the map var logic

* the tests are not valid + pyright fix

* darglint fix

* Add unit tests plus mix components

* pyi run

* rebase on main

* fix darglint

* testing different OS

* revert

* This should fix it. Right?

* Fix tests

* minor fn signature fix

* use ArgsFunctionOperation

* use destructured args and pass the tests

* fix remaining unit tests

* fix pyi files

* rebase on main

* move language regex on codeblock to markdown

* fix tests

---------

Co-authored-by: Khaleel Al-Adhami <[email protected]>

reflex/.templates/jinja/web/pages/custom_component.js.jinja2

@@ -8,20 +8,6 @@
 {% endfor %}
 
 export const {{component.name}} = memo(({ {{-component.props|join(", ")-}} }) => {
-{% if component.name == "CodeBlock" and "language" in component.props %}
-    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);
-        }
-      })();
-
-
-    }
-{% endif %}
     {% for hook in component.hooks %}
     {{ hook }}
     {% endfor %}

reflex/components/datadisplay/code.py

@@ -8,13 +8,14 @@
 from reflex.components.component import Component, ComponentNamespace
 from reflex.components.core.cond import color_mode_cond
 from reflex.components.lucide.icon import Icon
+from reflex.components.markdown.markdown import _LANGUAGE, MarkdownComponentMap
 from reflex.components.radix.themes.components.button import Button
 from reflex.components.radix.themes.layout.box import Box
 from reflex.constants.colors import Color
 from reflex.event import set_clipboard
 from reflex.style import Style
 from reflex.utils import console, format
-from reflex.utils.imports import ImportDict, ImportVar
+from reflex.utils.imports import ImportVar
 from reflex.vars.base import LiteralVar, Var, VarData
 
 LiteralCodeLanguage = Literal[
@@ -378,7 +379,7 @@ class Theme:
     setattr(Theme, theme_name, getattr(Theme, theme_name)._replace(_var_type=Theme))
 
 
-class CodeBlock(Component):
+class CodeBlock(Component, MarkdownComponentMap):
     """A code block."""
 
     library = "[email protected]"
@@ -417,39 +418,6 @@ class CodeBlock(Component):
     # A custom copy button to override the default one.
     copy_button: Optional[Union[bool, Component]] = None
 
-    def add_imports(self) -> ImportDict:
-        """Add imports for the CodeBlock component.
-
-        Returns:
-            The import dict.
-        """
-        imports_: ImportDict = {}
-
-        if (
-            self.language is not None
-            and (language_without_quotes := str(self.language).replace('"', ""))
-            in LiteralCodeLanguage.__args__  # type: ignore
-        ):
-            imports_[
-                f"react-syntax-highlighter/dist/cjs/languages/prism/{language_without_quotes}"
-            ] = [
-                ImportVar(
-                    tag=format.to_camel_case(language_without_quotes),
-                    is_default=True,
-                    install=False,
-                )
-            ]
-
-        return imports_
-
-    def _get_custom_code(self) -> Optional[str]:
-        if (
-            self.language is not None
-            and (language_without_quotes := str(self.language).replace('"', ""))
-            in LiteralCodeLanguage.__args__  # type: ignore
-        ):
-            return f"{self.alias}.registerLanguage('{language_without_quotes}', {format.to_camel_case(language_without_quotes)})"
-
     @classmethod
     def create(
         cls,
@@ -534,15 +502,55 @@ def _render(self):
 
         theme = self.theme
 
-        out.add_props(style=theme).remove_props("theme", "code").add_props(
-            children=self.code
+        out.add_props(style=theme).remove_props("theme", "code", "language").add_props(
+            children=self.code, language=_LANGUAGE
         )
 
         return out
 
     def _exclude_props(self) -> list[str]:
         return ["can_copy", "copy_button"]
 
+    @classmethod
+    def _get_language_registration_hook(cls) -> str:
+        """Get the hook to register the language.
+
+        Returns:
+            The hook to register the language.
+        """
+        return f"""
+ if ({str(_LANGUAGE)}) {{
+    (async () => {{
+      try {{
+        const module = await import(`react-syntax-highlighter/dist/cjs/languages/prism/${{{str(_LANGUAGE)}}}`);
+        SyntaxHighlighter.registerLanguage({str(_LANGUAGE)}, module.default);
+      }} catch (error) {{
+        console.error(`Error importing language module for ${{{str(_LANGUAGE)}}}:`, error);
+      }}
+    }})();
+  }}
+"""
+
+    @classmethod
+    def get_component_map_custom_code(cls) -> str:
+        """Get the custom code for the component.
+
+        Returns:
+            The custom code for the component.
+        """
+        return cls._get_language_registration_hook()
+
+    def add_hooks(self) -> list[str | Var]:
+        """Add hooks for the component.
+
+        Returns:
+            The hooks for the component.
+        """
+        return [
+            f"const {str(_LANGUAGE)} = {str(self.language)}",
+            self._get_language_registration_hook(),
+        ]
+
 
 class CodeblockNamespace(ComponentNamespace):
     """Namespace for the CodeBlock component."""

reflex/components/datadisplay/code.pyi

@@ -7,10 +7,10 @@ import dataclasses
 from typing import Any, ClassVar, Dict, Literal, Optional, Union, overload
 
 from reflex.components.component import Component, ComponentNamespace
+from reflex.components.markdown.markdown import MarkdownComponentMap
 from reflex.constants.colors import Color
 from reflex.event import BASE_STATE, EventType
 from reflex.style import Style
-from reflex.utils.imports import ImportDict
 from reflex.vars.base import Var
 
 LiteralCodeLanguage = Literal[
@@ -349,8 +349,7 @@ for theme_name in dir(Theme):
         continue
     setattr(Theme, theme_name, getattr(Theme, theme_name)._replace(_var_type=Theme))
 
-class CodeBlock(Component):
-    def add_imports(self) -> ImportDict: ...
+class CodeBlock(Component, MarkdownComponentMap):
     @overload
     @classmethod
     def create(  # type: ignore
@@ -984,6 +983,9 @@ class CodeBlock(Component):
         ...
 
     def add_style(self): ...
+    @classmethod
+    def get_component_map_custom_code(cls) -> str: ...
+    def add_hooks(self) -> list[str | Var]: ...
 
 class CodeblockNamespace(ComponentNamespace):
     themes = Theme

reflex/components/datadisplay/shiki_code_block.py

@@ -12,6 +12,7 @@
 from reflex.components.core.cond import color_mode_cond
 from reflex.components.el.elements.forms import Button
 from reflex.components.lucide.icon import Icon
+from reflex.components.markdown.markdown import MarkdownComponentMap
 from reflex.components.props import NoExtrasAllowedProps
 from reflex.components.radix.themes.layout.box import Box
 from reflex.event import run_script, set_clipboard
@@ -528,7 +529,7 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
 
-class ShikiCodeBlock(Component):
+class ShikiCodeBlock(Component, MarkdownComponentMap):
     """A Code block."""
 
     library = "/components/shiki/code"

reflex/components/datadisplay/shiki_code_block.pyi

@@ -7,6 +7,7 @@ from typing import Any, Dict, Literal, Optional, Union, overload
 
 from reflex.base import Base
 from reflex.components.component import Component, ComponentNamespace
+from reflex.components.markdown.markdown import MarkdownComponentMap
 from reflex.components.props import NoExtrasAllowedProps
 from reflex.event import BASE_STATE, EventType
 from reflex.style import Style
@@ -350,7 +351,7 @@ class ShikiJsTransformer(ShikiBaseTransformers):
     fns: list[FunctionStringVar]
     style: Optional[Style]
 
-class ShikiCodeBlock(Component):
+class ShikiCodeBlock(Component, MarkdownComponentMap):
     @overload
     @classmethod
     def create(  # type: ignore