rx.markdown custom styles for tags (#1416)

reflex/components/tags/tag.py

@@ -67,46 +67,48 @@ class Tag(Base):
         Raises:
             TypeError: If the prop is not a valid type.
         """
-        # Handle var props.
-        if isinstance(prop, Var):
-            if not prop.is_local or prop.is_string:
-                return str(prop)
-            if types._issubclass(prop.type_, str):
-                return format.json_dumps(prop.full_name)
-            prop = prop.full_name
-
-        # Handle event props.
-        elif isinstance(prop, EventChain):
-            if prop.full_control:
-                # Full control component events.
-                event = format.format_full_control_event(prop)
+        try:
+            # Handle var props.
+            if isinstance(prop, Var):
+                if not prop.is_local or prop.is_string:
+                    return str(prop)
+                if types._issubclass(prop.type_, str):
+                    return format.json_dumps(prop.full_name)
+                prop = prop.full_name
+
+            # Handle event props.
+            elif isinstance(prop, EventChain):
+                if prop.full_control:
+                    # Full control component events.
+                    event = format.format_full_control_event(prop)
+                else:
+                    # All other events.
+                    chain = ",".join(
+                        [format.format_event(event) for event in prop.events]
+                    )
+                    event = f"Event([{chain}], {EVENT_ARG})"
+                prop = f"{EVENT_ARG} => {event}"
+
+            # Handle other types.
+            elif isinstance(prop, str):
+                if format.is_wrapped(prop, "{"):
+                    return prop
+                return format.json_dumps(prop)
+
+            elif isinstance(prop, Figure):
+                prop = json.loads(to_json(prop))["data"]  # type: ignore
+
+            # For dictionaries, convert any properties to strings.
+            elif isinstance(prop, dict):
+                prop = format.format_dict(prop)
+
             else:
-                # All other events.
-                chain = ",".join([format.format_event(event) for event in prop.events])
-                event = f"Event([{chain}], {EVENT_ARG})"
-            prop = f"{EVENT_ARG} => {event}"
-
-        # Handle other types.
-        elif isinstance(prop, str):
-            if format.is_wrapped(prop, "{"):
-                return prop
-            return format.json_dumps(prop)
-
-        elif isinstance(prop, Figure):
-            prop = json.loads(to_json(prop))["data"]  # type: ignore
-
-        # For dictionaries, convert any properties to strings.
-        elif isinstance(prop, dict):
-            prop = format.format_dict(prop)
-
-        else:
-            # Dump the prop as JSON.
-            try:
+                # Dump the prop as JSON.
                 prop = format.json_dumps(prop)
-            except TypeError as e:
-                raise TypeError(
-                    f"Could not format prop: {prop} of type {type(prop)}"
-                ) from e
+        except TypeError as e:
+            raise TypeError(
+                f"Could not format prop: {prop} of type {type(prop)}"
+            ) from e
 
         # Wrap the variable in braces.
         assert isinstance(prop, str), "The prop must be a string."

reflex/components/typography/heading.py

@@ -9,5 +9,8 @@ class Heading(ChakraComponent):
 
     tag = "Heading"
 
+    # Override the tag. The default tag is `<h2>`.
+    as_: Var[str]
+
     # "4xl" | "3xl" | "2xl" | "xl" | "lg" | "md" | "sm" | "xs"
     size: Var[str]

reflex/components/typography/markdown.py

@@ -1,12 +1,33 @@
 """Markdown component."""
 
 import textwrap
-from typing import List, Union
+from typing import Callable, Dict, List, Union
 
+from reflex.compiler import utils
 from reflex.components.component import Component
+from reflex.components.datadisplay.list import ListItem, OrderedList, UnorderedList
+from reflex.components.navigation import Link
+from reflex.components.typography.heading import Heading
+from reflex.components.typography.text import Text
+from reflex.style import Style
 from reflex.utils import types
 from reflex.vars import BaseVar, ImportVar, Var
 
+# Mapping from markdown tags to components.
+components_by_tag: Dict[str, Callable] = {
+    "h1": Heading,
+    "h2": Heading,
+    "h3": Heading,
+    "h4": Heading,
+    "h5": Heading,
+    "h6": Heading,
+    "p": Text,
+    "ul": UnorderedList,
+    "ol": OrderedList,
+    "li": ListItem,
+    "a": Link,
+}
+
 
 class Markdown(Component):
     """A markdown component."""
@@ -17,6 +38,37 @@ class Markdown(Component):
 
     is_default = True
 
+    # Custom defined styles for the markdown elements.
+    custom_styles: Dict[str, Style] = {
+        k: Style(v)
+        for k, v in {
+            "h1": {
+                "as_": "h1",
+                "size": "2xl",
+            },
+            "h2": {
+                "as_": "h2",
+                "size": "xl",
+            },
+            "h3": {
+                "as_": "h3",
+                "size": "lg",
+            },
+            "h4": {
+                "as_": "h4",
+                "size": "md",
+            },
+            "h5": {
+                "as_": "h5",
+                "size": "sm",
+            },
+            "h6": {
+                "as_": "h6",
+                "size": "xs",
+            },
+        }.items()
+    }
+
     @classmethod
     def create(cls, *children, **props) -> Component:
         """Create a markdown component.
@@ -39,47 +91,45 @@ class Markdown(Component):
         return super().create(src, **props)
 
     def _get_imports(self):
+        from reflex.components.datadisplay.code import Code, CodeBlock
+
         imports = super()._get_imports()
-        imports["@chakra-ui/react"] = {
-            ImportVar(tag="Heading"),
-            ImportVar(tag="Code"),
-            ImportVar(tag="Text"),
-            ImportVar(tag="Link"),
-            ImportVar(tag="UnorderedList"),
-            ImportVar(tag="OrderedList"),
-            ImportVar(tag="ListItem"),
-        }
-        imports["react-syntax-highlighter"] = {ImportVar(tag="Prism", is_default=True)}
+
+        # Special markdown imports.
         imports["remark-math"] = {ImportVar(tag="remarkMath", is_default=True)}
         imports["remark-gfm"] = {ImportVar(tag="remarkGfm", is_default=True)}
         imports["rehype-katex"] = {ImportVar(tag="rehypeKatex", is_default=True)}
         imports["rehype-raw"] = {ImportVar(tag="rehypeRaw", is_default=True)}
         imports[""] = {ImportVar(tag="katex/dist/katex.min.css")}
+
+        # Get the imports for each component.
+        for component in components_by_tag.values():
+            imports = utils.merge_imports(imports, component()._get_imports())
+
+        # Get the imports for the code components.
+        imports = utils.merge_imports(
+            imports, CodeBlock.create(theme="light")._get_imports()
+        )
+        imports = utils.merge_imports(imports, Code.create()._get_imports())
         return imports
 
     def _render(self):
-        return (
-            super()
-            ._render()
-            .add_props(
-                components={
-                    "h1": "{({node, ...props}) => <Heading size='2xl' paddingY='0.5em' {...props} />}",
-                    "h2": "{({node, ...props}) => <Heading size='xl' paddingY='0.5em' {...props} />}",
-                    "h3": "{({node, ...props}) => <Heading size='lg' paddingY='0.5em' {...props} />}",
-                    "h4": "{({node, ...props}) => <Heading size='sm' paddingY='0.5em' {...props} />}",
-                    "h5": "{({node, ...props}) => <Heading size='xs' paddingY='0.5em' {...props} />}",
-                    "ul": "{UnorderedList}",
-                    "ol": "{OrderedList}",
-                    "li": "{ListItem}",
-                    "p": "{({node, ...props}) => <Text paddingY='0.5em' {...props} />}",
-                    "a": "{Link}",
-                    "code": """{({node, inline, className, children, ...props}) =>
+        from reflex.components.tags.tag import Tag
+
+        components = {
+            tag: f"{{({{node, ...props}}) => <{(component().tag)} {{...props}} {''.join(Tag(name='', props=Style(self.custom_styles.get(tag, {}))).format_props())} />}}"
+            for tag, component in components_by_tag.items()
+        }
+        components[
+            "code"
+        ] = """{({node, inline, className, children, ...props}) =>
                     {
         const match = (className || '').match(/language-(?<lang>.*)/);
         return !inline ? (
           <Prism
             children={String(children).replace(/\n$/, '')}
             language={match ? match[1] : ''}
+            theme={light}
             {...props}
           />
         ) : (
@@ -88,12 +138,18 @@ class Markdown(Component):
           </Code>
         );
       }}""".replace(
-                        "\n", " "
-                    ),
-                },
+            "\n", " "
+        )
+
+        return (
+            super()
+            ._render()
+            .add_props(
+                components=components,
                 remark_plugins=BaseVar(name="[remarkMath, remarkGfm]", type_=List[str]),
                 rehype_plugins=BaseVar(
                     name="[rehypeKatex, rehypeRaw]", type_=List[str]
                 ),
             )
+            .remove_props("custom_components")
         )