improve behavior on missing language with markdown code blocks (#4750)

* improve behavior on missing language with markdown code blocks

* special case on literal var

* fix tests

* missing f

* remove extra throw

reflex/components/datadisplay/code.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import dataclasses
+import typing
 from typing import ClassVar, Dict, Literal, Optional, Union
 
 from reflex.components.component import Component, ComponentNamespace
@@ -503,7 +504,7 @@ class CodeBlock(Component, MarkdownComponentMap):
         return ["can_copy", "copy_button"]
 
     @classmethod
-    def _get_language_registration_hook(cls, language_var: Var = _LANGUAGE) -> str:
+    def _get_language_registration_hook(cls, language_var: Var = _LANGUAGE) -> Var:
         """Get the hook to register the language.
 
         Args:
@@ -514,21 +515,46 @@ class CodeBlock(Component, MarkdownComponentMap):
         Returns:
             The hook to register the language.
         """
-        return f"""
- if ({language_var!s}) {{
-    (async () => {{
-      try {{
+        language_in_there = Var.create(typing.get_args(LiteralCodeLanguage)).contains(
+            language_var
+        )
+        async_load = f"""
+(async () => {{
+    try {{
         const module = await import(`react-syntax-highlighter/dist/cjs/languages/prism/${{{language_var!s}}}`);
         SyntaxHighlighter.registerLanguage({language_var!s}, module.default);
-      }} catch (error) {{
-        console.error(`Error importing language module for ${{{language_var!s}}}:`, error);
-      }}
-    }})();
+    }} catch (error) {{
+        console.error(`Language ${{{language_var!s}}} is not supported for code blocks inside of markdown: `, error);
+    }}
+}})();
+"""
+        return Var(
+            f"""
+ if ({language_var!s}) {{
+    if (!{language_in_there!s}) {{
+        console.warn(`Language \\`${{{language_var!s}}}\\` is not supported for code blocks inside of markdown.`);
+        {language_var!s} = '';
+    }} else {{
+        {async_load!s}
+    }}
   }}
 """
+            if not isinstance(language_var, LiteralVar)
+            else f"""
+if ({language_var!s}) {{
+    {async_load!s}
+}}""",
+            _var_data=VarData(
+                imports={
+                    cls.__fields__["library"].default: [
+                        ImportVar(tag="PrismAsyncLight", alias="SyntaxHighlighter")
+                    ]
+                },
+            ),
+        )
 
     @classmethod
-    def get_component_map_custom_code(cls) -> str:
+    def get_component_map_custom_code(cls) -> Var:
         """Get the custom code for the component.
 
         Returns:

reflex/components/datadisplay/code.pyi

@@ -984,7 +984,7 @@ class CodeBlock(Component, MarkdownComponentMap):
 
     def add_style(self): ...
     @classmethod
-    def get_component_map_custom_code(cls) -> str: ...
+    def get_component_map_custom_code(cls) -> Var: ...
     def add_hooks(self) -> list[str | Var]: ...
 
 class CodeblockNamespace(ComponentNamespace):

reflex/components/markdown/markdown.py

@@ -12,7 +12,7 @@ from reflex.components.component import BaseComponent, Component, CustomComponen
 from reflex.components.tags.tag import Tag
 from reflex.utils import types
 from reflex.utils.imports import ImportDict, ImportVar
-from reflex.vars.base import LiteralVar, Var
+from reflex.vars.base import LiteralVar, Var, VarData
 from reflex.vars.function import ARRAY_ISARRAY, ArgsFunctionOperation, DestructuredArg
 from reflex.vars.number import ternary_operation
 
@@ -83,13 +83,13 @@ class MarkdownComponentMap:
     _explicit_return: bool = dataclasses.field(default=False)
 
     @classmethod
-    def get_component_map_custom_code(cls) -> str:
+    def get_component_map_custom_code(cls) -> Var:
         """Get the custom code for the component map.
 
         Returns:
             The custom code for the component map.
         """
-        return ""
+        return Var("")
 
     @classmethod
     def create_map_fn_var(
@@ -97,6 +97,7 @@ class MarkdownComponentMap:
         fn_body: Var | None = None,
         fn_args: Sequence[str] | None = None,
         explicit_return: bool | None = None,
+        var_data: VarData | None = None,
     ) -> Var:
         """Create a function Var for the component map.
 
@@ -104,6 +105,7 @@ class MarkdownComponentMap:
             fn_body: The formatted component as a string.
             fn_args: The function arguments.
             explicit_return: Whether to use explicit return syntax.
+            var_data: The var data for the function.
 
         Returns:
             The function Var for the component map.
@@ -116,6 +118,7 @@ class MarkdownComponentMap:
             args_names=(DestructuredArg(fields=tuple(fn_args)),),
             return_expr=fn_body,
             explicit_return=explicit_return,
+            _var_data=var_data,
         )
 
     @classmethod
@@ -239,6 +242,15 @@ class Markdown(Component):
                 component(_MOCK_ARG)._get_all_imports()
                 for component in self.component_map.values()
             ],
+            *(
+                [inline_code_var_data.old_school_imports()]
+                if (
+                    inline_code_var_data
+                    := self._get_inline_code_fn_var()._get_all_var_data()
+                )
+                is not None
+                else []
+            ),
         ]
 
     def _get_tag_map_fn_var(self, tag: str) -> Var:
@@ -278,12 +290,20 @@ class Markdown(Component):
             self._get_map_fn_custom_code_from_children(self.get_component("code"))
         )
 
-        codeblock_custom_code = "\n".join(custom_code_list)
+        var_data = VarData.merge(
+            *[
+                code._get_all_var_data()
+                for code in custom_code_list
+                if isinstance(code, Var)
+            ]
+        )
+
+        codeblock_custom_code = "\n".join(map(str, custom_code_list))
 
         # Format the code to handle inline and block code.
         formatted_code = f"""
 const match = (className || '').match(/language-(?<lang>.*)/);
-const {_LANGUAGE!s} = match ? match[1] : '';
+let {_LANGUAGE!s} = match ? match[1] : '';
 {codeblock_custom_code};
             return inline ? (
                 {self.format_component("code")}
@@ -302,6 +322,7 @@ const {_LANGUAGE!s} = match ? match[1] : '';
             ),
             fn_body=Var(_js_expr=formatted_code),
             explicit_return=True,
+            var_data=var_data,
         )
 
     def get_component(self, tag: str, **props) -> Component:
@@ -381,7 +402,7 @@ const {_LANGUAGE!s} = match ? match[1] : '';
 
     def _get_map_fn_custom_code_from_children(
         self, component: BaseComponent
-    ) -> list[str]:
+    ) -> list[str | Var]:
         """Recursively get markdown custom code from children components.
 
         Args:
@@ -390,7 +411,7 @@ const {_LANGUAGE!s} = match ? match[1] : '';
         Returns:
             A list of markdown custom code strings.
         """
-        custom_code_list = []
+        custom_code_list: list[str | Var] = []
         if isinstance(component, MarkdownComponentMap):
             custom_code_list.append(component.get_component_map_custom_code())
 

reflex/components/markdown/markdown.pyi

@@ -11,7 +11,7 @@ from reflex.components.component import Component
 from reflex.event import EventType
 from reflex.style import Style
 from reflex.utils.imports import ImportDict
-from reflex.vars.base import LiteralVar, Var
+from reflex.vars.base import LiteralVar, Var, VarData
 
 _CHILDREN = Var(_js_expr="children", _var_type=str)
 _PROPS = Var(_js_expr="...props")
@@ -32,13 +32,14 @@ def get_base_component_map() -> dict[str, Callable]: ...
 @dataclasses.dataclass()
 class MarkdownComponentMap:
     @classmethod
-    def get_component_map_custom_code(cls) -> str: ...
+    def get_component_map_custom_code(cls) -> Var: ...
     @classmethod
     def create_map_fn_var(
         cls,
         fn_body: Var | None = None,
         fn_args: Sequence[str] | None = None,
         explicit_return: bool | None = None,
+        var_data: VarData | None = None,
     ) -> Var: ...
     @classmethod
     def get_fn_args(cls) -> Sequence[str]: ...