Memoize markdown component_map (#2219)

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

@@ -3,6 +3,10 @@
 {% block export %}
 {% for component in components %}
 
+{% for custom_code in component.custom_code %}
+{{custom_code}}
+{% endfor %}
+
 export const {{component.name}} = memo(({ {{-component.props|join(", ")-}} }) => {
 {% if component.name == "CodeBlock" and "language" in component.props %}
     if (language) {

reflex/compiler/utils.py

@@ -254,6 +254,7 @@ def compile_custom_component(
             "name": component.tag,
             "props": props,
             "render": render.render(),
+            "custom_code": render.get_custom_code(),
         },
         imports,
     )

reflex/components/datadisplay/code.py

@@ -1,5 +1,5 @@
 """A code component."""
-
+import re
 from typing import Dict, Literal, Optional, Union
 
 from reflex.components.component import Component
@@ -362,6 +362,9 @@ class CodeBlock(Component):
     # The language to use.
     language: Var[LiteralCodeLanguage] = "python"  # type: ignore
 
+    # The code to display.
+    code: Var[str]
+
     # If this is enabled line numbers will be shown next to the code block.
     show_line_numbers: Var[bool]
 
@@ -379,16 +382,21 @@ class CodeBlock(Component):
 
     def _get_imports(self) -> imports.ImportDict:
         merged_imports = super()._get_imports()
+        # Get all themes from a cond literal
+        themes = re.findall(r"`(.*?)`", self.theme._var_name)
+        if not themes:
+            themes = [self.theme._var_name]
         merged_imports = imports.merge_imports(
             merged_imports,
             {
-                f"react-syntax-highlighter/dist/cjs/styles/prism/{self.theme._var_name}": {
+                f"react-syntax-highlighter/dist/cjs/styles/prism/{theme}": {
                     ImportVar(
-                        tag=format.to_camel_case(self.theme._var_name),
+                        tag=format.to_camel_case(theme),
                         is_default=True,
                         install=False,
                     )
                 }
+                for theme in themes
             },
         )
         if (
@@ -440,7 +448,7 @@ class CodeBlock(Component):
 
         # react-syntax-highlighter doesnt have an explicit "light" or "dark" theme so we use one-light and one-dark
         # themes respectively to ensure code compatibility.
-        if "theme" in props:
+        if "theme" in props and not isinstance(props["theme"], Var):
             props["theme"] = (
                 "one-light"
                 if props["theme"] == "light"
@@ -469,9 +477,14 @@ class CodeBlock(Component):
             if key not in cls.get_fields():
                 custom_style[key] = value
 
+        # Carry the children (code) via props
+        if children:
+            props["code"] = children[0]
+            if not isinstance(props["code"], Var):
+                props["code"] = Var.create(props["code"], _var_is_string=True)
+
         # Create the component.
         code_block = super().create(
-            *children,
             **props,
             custom_style=Style(custom_style),
         )
@@ -486,11 +499,15 @@ class CodeBlock(Component):
 
     def _render(self):
         out = super()._render()
+        predicate, qmark, value = self.theme._var_name.partition("?")
         out.add_props(
             style=Var.create(
-                format.to_camel_case(self.theme._var_name), _var_is_local=False
+                format.to_camel_case(f"{predicate}{qmark}{value.replace('`', '')}"),
+                _var_is_local=False,
             )
-        ).remove_props("theme")
+        ).remove_props("theme", "code")
+        if self.code is not None:
+            out.special_props.add(Var.create_safe(f"children={str(self.code)}"))
         return out
 
 

reflex/components/datadisplay/code.pyi

@@ -7,6 +7,7 @@ from typing import Any, Dict, Literal, Optional, Union, overload
 from reflex.vars import Var, BaseVar, ComputedVar
 from reflex.event import EventChain, EventHandler, EventSpec
 from reflex.style import Style
+import re
 from typing import Dict, Literal, Optional, Union
 from reflex.components.component import Component
 from reflex.components.forms import Button
@@ -1024,6 +1025,7 @@ class CodeBlock(Component):
                 ],
             ]
         ] = None,
+        code: Optional[Union[Var[str], str]] = None,
         show_line_numbers: Optional[Union[Var[bool], bool]] = None,
         starting_line_number: Optional[Union[Var[int], int]] = None,
         wrap_long_lines: Optional[Union[Var[bool], bool]] = None,
@@ -1090,6 +1092,7 @@ class CodeBlock(Component):
             copy_button: A custom copy button to override the default one.
             theme: The theme to use ("light" or "dark").
             language: The language to use.
+            code: The code to display.
             show_line_numbers: If this is enabled line numbers will be shown next to the code block.
             starting_line_number: The starting line number to use.
             wrap_long_lines: Whether to wrap long lines.

reflex/components/typography/markdown.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import textwrap
+from hashlib import md5
 from typing import Any, Callable, Dict, Union
 
 from reflex.compiler import utils
@@ -67,8 +68,8 @@ def get_base_component_map() -> dict[str, Callable]:
         "li": lambda value: ListItem.create(value, margin_y="0.5em"),
         "a": lambda value: Link.create(value),
         "code": lambda value: Code.create(value),
-        "codeblock": lambda *_, **props: CodeBlock.create(
+        "codeblock": lambda value, **props: CodeBlock.create(
-            theme="light", margin_y="1em", **props
+            value, theme="light", margin_y="1em", **props
         ),
     }
 
@@ -232,14 +233,14 @@ class Markdown(Component):
             The formatted component map.
         """
         components = {
-            tag: f"{{({{{_CHILDREN._var_name}, {_PROPS._var_name}}}) => {self.format_component(tag)}}}"
+            tag: f"{{({{node, {_CHILDREN._var_name}, {_PROPS._var_name}}}) => {self.format_component(tag)}}}"
             for tag in self.component_map
         }
 
         # Separate out inline code and code blocks.
         components[
             "code"
-        ] = f"""{{({{inline, className, {_CHILDREN._var_name}, {_PROPS._var_name}}}) => {{
+        ] = f"""{{({{node, inline, className, {_CHILDREN._var_name}, {_PROPS._var_name}}}) => {{
     const match = (className || '').match(/language-(?<lang>.*)/);
     const language = match ? match[1] : '';
     if (language) {{
@@ -255,7 +256,7 @@ class Markdown(Component):
     return inline ? (
         {self.format_component("code")}
     ) : (
-        {self.format_component("codeblock", language=Var.create_safe("language", _var_is_local=False), children=Var.create_safe("String(children)", _var_is_local=False))}
+        {self.format_component("codeblock", language=Var.create_safe("language", _var_is_local=False))}
     );
       }}}}""".replace(
             "\n", " "
@@ -263,14 +264,41 @@ class Markdown(Component):
 
         return components
 
+    def _component_map_hash(self) -> str:
+        return md5(str(self.component_map).encode()).hexdigest()
+
+    def _get_component_map_name(self) -> str:
+        return f"ComponentMap_{self._component_map_hash()}"
+
+    def _get_custom_code(self) -> str | None:
+        hooks = set()
+        for component in self.component_map.values():
+            hooks |= component(_MOCK_ARG).get_hooks()
+        formatted_hooks = "\n".join(hooks)
+        return f"""
+        function {self._get_component_map_name()} () {{
+            {formatted_hooks}
+            return (
+                {str(Var.create(self.format_component_map()))}
+            )
+        }}
+        """
+
     def _render(self) -> Tag:
-        return (
+        tag = (
             super()
             ._render()
             .add_props(
-                components=self.format_component_map(),
                 remark_plugins=_REMARK_PLUGINS,
                 rehype_plugins=_REHYPE_PLUGINS,
             )
             .remove_props("componentMap")
         )
+        tag.special_props.add(
+            Var.create_safe(
+                f"components={{{self._get_component_map_name()}()}}",
+                _var_is_local=True,
+                _var_is_string=False,
+            ),
+        )
+        return tag

reflex/components/typography/markdown.pyi

@@ -8,6 +8,7 @@ from reflex.vars import Var, BaseVar, ComputedVar
 from reflex.event import EventChain, EventHandler, EventSpec
 from reflex.style import Style
 import textwrap
+from hashlib import md5
 from typing import Any, Callable, Dict, Union
 from reflex.compiler import utils
 from reflex.components.component import Component, CustomComponent