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

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

@@ -379,6 +379,8 @@ for theme_name in dir(Theme):
     setattr(Theme, theme_name, getattr(Theme, theme_name)._replace(_var_type=Theme))
 
 
+LANGUAGE_VAR = Var(_js_expr="__language")
+
 class CodeBlock(Component, MarkdownComponentMapMixin):
     """A code block."""
 
@@ -426,30 +428,13 @@ class CodeBlock(Component, MarkdownComponentMapMixin):
         """
         imports_: ImportDict = {}
 
-        if (
+    # def _get_custom_code(self) -> Optional[str]:
-            self.language is not None
+    #     if (
-            and (language_without_quotes := str(self.language).replace('"', ""))
+    #         self.language is not None
-            in LiteralCodeLanguage.__args__  # type: ignore
+    #         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}"
+    #         return f"{self.alias}.registerLanguage('{language_without_quotes}', {format.to_camel_case(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(
@@ -535,8 +520,9 @@ class CodeBlock(Component, MarkdownComponentMapMixin):
 
         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_VAR
         )
 
         return out
@@ -567,6 +553,29 @@ const match = (className || '').match(/language-(?<lang>.*)/);
         """
 
 
+    def add_hooks(self) -> list[str | Var]:
+        """Add hooks for the component.
+
+        Returns:
+            The hooks for the component.
+        """
+        return [
+            f"const {str(LANGUAGE_VAR)} = {str(self.language)}",
+            f"""
+ if ({str(LANGUAGE_VAR)}) {{
+    (async () => {{
+      try {{
+        const module = await import(`react-syntax-highlighter/dist/cjs/languages/prism/${{{str(LANGUAGE_VAR)}}}`);
+        SyntaxHighlighter.registerLanguage({str(LANGUAGE_VAR)}, module.default);
+      }} catch (error) {{
+        console.error(`Error importing language module for ${{{str(LANGUAGE_VAR)}}}:`, error);
+      }}
+    }})();
+  }}                
+            """
+        ]
+
+
 class CodeblockNamespace(ComponentNamespace):
     """Namespace for the CodeBlock component."""
 

reflex/components/markdown/markdown.py

@@ -257,8 +257,8 @@ 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,
+        custom_code_list = self._get_custom_code_from_children(codeblock_component)
-                                                                                               "get_component_map_custom_code") else ""
+        codeblock_custom_code = "\n".join(custom_code_list)
         # Separate out inline code and code blocks.
         components["code"] = Var(
             _js_expr=f"""(({{node, inline, className, {_CHILDREN._js_expr}, {_PROPS._js_expr}}}) => {{
@@ -273,6 +273,26 @@ class Markdown(Component):
 
         return components
 
+    def _get_custom_code_from_children(self, component) -> list[str]:
+        """Recursively get markdown custom code from children components.
+
+        Args:
+            component: The component to check for custom code.
+
+        Returns:
+            A list of markdown custom code strings.
+        """
+        custom_code_list = []
+        if hasattr(component, "get_component_map_custom_code"):
+            custom_code_list.append(component.get_component_map_custom_code())
+
+        if isinstance(component, CustomComponent):
+            custom_code_list.extend(self._get_custom_code_from_children(component.component_fn(*component.get_prop_vars())))
+        else:
+            for child in component.children:
+                custom_code_list.extend(self._get_custom_code_from_children(child))
+
+        return custom_code_list
     @staticmethod
     def _component_map_hash(component_map) -> str:
         inp = str(
@@ -284,12 +304,12 @@ class Markdown(Component):
         return f"ComponentMap_{self.component_map_hash}"
 
     def _get_custom_code(self) -> str | None:
-        hooks = set()
+        hooks = {}
         for _component in self.component_map.values():
             comp = _component(_MOCK_ARG)
             hooks.update(comp._get_all_hooks_internal())
             hooks.update(comp._get_all_hooks())
-        formatted_hooks = "\n".join(hooks)
+        formatted_hooks = "\n".join(hooks.keys())
         return f"""
         function {self._get_component_map_name()} () {{
             {formatted_hooks}