From 626357ed87b0fbece46d5b7b5f2db6dee26640db Mon Sep 17 00:00:00 2001 From: Masen Furer Date: Mon, 27 Nov 2023 18:09:41 -0800 Subject: [PATCH] Memoize markdown component_map (#2219) --- .../web/pages/custom_component.js.jinja2 | 4 ++ reflex/compiler/utils.py | 1 + reflex/components/datadisplay/code.py | 31 ++++++++++---- reflex/components/datadisplay/code.pyi | 3 ++ reflex/components/typography/markdown.py | 42 +++++++++++++++---- reflex/components/typography/markdown.pyi | 1 + 6 files changed, 68 insertions(+), 14 deletions(-) diff --git a/reflex/.templates/jinja/web/pages/custom_component.js.jinja2 b/reflex/.templates/jinja/web/pages/custom_component.js.jinja2 index 643651486..6b1438c62 100644 --- a/reflex/.templates/jinja/web/pages/custom_component.js.jinja2 +++ b/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) { diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index 074593b4b..3a866051a 100644 --- a/reflex/compiler/utils.py +++ b/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, ) diff --git a/reflex/components/datadisplay/code.py b/reflex/components/datadisplay/code.py index 0c796cd08..046ad4e70 100644 --- a/reflex/components/datadisplay/code.py +++ b/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 diff --git a/reflex/components/datadisplay/code.pyi b/reflex/components/datadisplay/code.pyi index 797d30958..a35f22512 100644 --- a/reflex/components/datadisplay/code.pyi +++ b/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. diff --git a/reflex/components/typography/markdown.py b/reflex/components/typography/markdown.py index 5ebd595b7..fc58beef9 100644 --- a/reflex/components/typography/markdown.py +++ b/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( - theme="light", margin_y="1em", **props + "codeblock": lambda value, **props: CodeBlock.create( + 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-(?.*)/); 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 diff --git a/reflex/components/typography/markdown.pyi b/reflex/components/typography/markdown.pyi index e5f9d09f7..3a0fb43a8 100644 --- a/reflex/components/typography/markdown.pyi +++ b/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