From 3113aecb307c13ee05b123f0c26dad9b7c0078cc Mon Sep 17 00:00:00 2001
From: Nikhil Rao <nikhil@pynecone.io>
Date: Wed, 20 Sep 2023 11:40:25 -0700
Subject: [PATCH] Support custom styling for code in markdown (#1844)

---
 reflex/components/tags/tag.py            | 11 +-----
 reflex/components/typography/markdown.py | 44 +++++++++++++++---------
 reflex/utils/format.py                   | 18 ++++++++++
 3 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/reflex/components/tags/tag.py b/reflex/components/tags/tag.py
index 855b130566..7c63443696 100644
--- a/reflex/components/tags/tag.py
+++ b/reflex/components/tags/tag.py
@@ -51,16 +51,7 @@ def format_props(self) -> List:
         Returns:
             The formatted props list.
         """
-        # If there are no props, return an empty string.
-        if len(self.props) == 0:
-            return []
-
-        # Format all the props.
-        return [
-            f"{name}={format.format_prop(prop)}"
-            for name, prop in sorted(self.props.items())
-            if prop is not None
-        ] + [str(prop) for prop in self.special_props]
+        return format.format_props(*self.special_props, **self.props)
 
     def add_props(self, **kwargs: Optional[Any]) -> Tag:
         """Add props to the tag.
diff --git a/reflex/components/typography/markdown.py b/reflex/components/typography/markdown.py
index baf0824d47..935f4885e1 100644
--- a/reflex/components/typography/markdown.py
+++ b/reflex/components/typography/markdown.py
@@ -91,6 +91,7 @@ def create(cls, *children, **props) -> Component:
         return super().create(src, **props)
 
     def _get_imports(self):
+        # Import here to avoid circular imports.
         from reflex.components.datadisplay.code import Code, CodeBlock
 
         imports = super()._get_imports()
@@ -118,30 +119,39 @@ def _get_imports(self):
         return imports
 
     def _render(self):
+        # Import here to avoid circular imports.
+        from reflex.components.datadisplay.code import Code, CodeBlock
         from reflex.components.tags.tag import Tag
 
+        def format_props(tag):
+            return "".join(
+                Tag(
+                    name="", props=Style(self.custom_styles.get(tag, {}))
+                ).format_props()
+            )
+
         components = {
-            tag: f"{{({{node, ...props}}) => <{(component().tag)} {{...props}} {''.join(Tag(name='', props=Style(self.custom_styles.get(tag, {}))).format_props())} />}}"
+            tag: f"{{({{node, ...props}}) => <{(component().tag)} {{...props}} {format_props(tag)} />}}"
             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}
-          />
-        ) : (
-          <Code {...props}>
-            {children}
-          </Code>
-        );
-      }}""".replace(
+        ] = f"""{{({{node, inline, className, children, ...props}}) => {{
+    const match = (className || '').match(/language-(?<lang>.*)/);
+    return !inline ? (
+        <{CodeBlock().tag}
+        children={{String(children).replace(/\n$/, '')}}
+        language={{match ? match[1] : ''}}
+        style={{light}}
+        {{...props}}
+        {format_props("pre")}
+        />
+    ) : (
+        <{Code.create().tag} {{...props}} {format_props("code")}>
+        {{children}}
+        </{Code.create().tag}>
+    );
+      }}}}""".replace(
             "\n", " "
         )
 
diff --git a/reflex/utils/format.py b/reflex/utils/format.py
index 21c08cb95d..e4e79ca669 100644
--- a/reflex/utils/format.py
+++ b/reflex/utils/format.py
@@ -327,6 +327,24 @@ def format_prop(
     return wrap(prop, "{", check_first=False)
 
 
+def format_props(*single_props, **key_value_props) -> list[str]:
+    """Format the tag's props.
+
+    Args:
+        single_props: Props that are not key-value pairs.
+        key_value_props: Props that are key-value pairs.
+
+    Returns:
+        The formatted props list.
+    """
+    # Format all the props.
+    return [
+        f"{name}={format_prop(prop)}"
+        for name, prop in sorted(key_value_props.items())
+        if prop is not None
+    ] + [str(prop) for prop in sorted(single_props)]
+
+
 def get_event_handler_parts(handler: EventHandler) -> tuple[str, str]:
     """Get the state and function name of an event handler.