code.py

"""A code component."""

from __future__ import annotations

import dataclasses
from collections.abc import Iterator
from typing import ClassVar, Literal

from reflex_base.components.component import Component, ComponentNamespace, field
from reflex_base.constants.colors import Color
from reflex_base.event import set_clipboard
from reflex_base.style import Style
from reflex_base.utils import format
from reflex_base.utils.imports import ImportVar
from reflex_base.vars.base import LiteralVar, Var, VarData
from reflex_components_core.core.cond import color_mode_cond
from reflex_components_core.core.markdown_component_map import MarkdownComponentMap
from reflex_components_radix.themes.components.button import Button
from reflex_components_radix.themes.layout.box import Box

LiteralCodeLanguage = Literal[
    "abap",
    "abnf",
    "actionscript",
    "ada",
    "agda",
    "al",
    "antlr4",
    "apacheconf",
    "apex",
    "apl",
    "applescript",
    "aql",
    "arduino",
    "arff",
    "asciidoc",
    "asm6502",
    "asmatmel",
    "aspnet",
    "autohotkey",
    "autoit",
    "avisynth",
    "avro-idl",
    "bash",
    "basic",
    "batch",
    "bbcode",
    "bicep",
    "birb",
    "bison",
    "bnf",
    "brainfuck",
    "brightscript",
    "bro",
    "bsl",
    "c",
    "cfscript",
    "chaiscript",
    "cil",
    "clike",
    "clojure",
    "cmake",
    "cobol",
    "coffeescript",
    "concurnas",
    "coq",
    "core",
    "cpp",
    "crystal",
    "csharp",
    "cshtml",
    "csp",
    "css",
    "css-extras",
    "csv",
    "cypher",
    "d",
    "dart",
    "dataweave",
    "dax",
    "dhall",
    "diff",
    "django",
    "dns-zone-file",
    "docker",
    "dot",
    "ebnf",
    "editorconfig",
    "eiffel",
    "ejs",
    "elixir",
    "elm",
    "erb",
    "erlang",
    "etlua",
    "excel-formula",
    "factor",
    "false",
    "firestore-security-rules",
    "flow",
    "fortran",
    "fsharp",
    "ftl",
    "gap",
    "gcode",
    "gdscript",
    "gedcom",
    "gherkin",
    "git",
    "glsl",
    "gml",
    "gn",
    "go",
    "go-module",
    "graphql",
    "groovy",
    "haml",
    "handlebars",
    "haskell",
    "haxe",
    "hcl",
    "hlsl",
    "hoon",
    "hpkp",
    "hsts",
    "http",
    "ichigojam",
    "icon",
    "icu-message-format",
    "idris",
    "iecst",
    "ignore",
    "index",
    "inform7",
    "ini",
    "io",
    "j",
    "java",
    "javadoc",
    "javadoclike",
    "javascript",
    "javastacktrace",
    "jexl",
    "jolie",
    "jq",
    "js-extras",
    "js-templates",
    "jsdoc",
    "json",
    "json5",
    "jsonp",
    "jsstacktrace",
    "jsx",
    "julia",
    "keepalived",
    "keyman",
    "kotlin",
    "kumir",
    "kusto",
    "latex",
    "latte",
    "less",
    "lilypond",
    "liquid",
    "lisp",
    "livescript",
    "llvm",
    "log",
    "lolcode",
    "lua",
    "magma",
    "makefile",
    "markdown",
    "markup",
    "markup-templating",
    "matlab",
    "maxscript",
    "mel",
    "mermaid",
    "mizar",
    "mongodb",
    "monkey",
    "moonscript",
    "n1ql",
    "n4js",
    "nand2tetris-hdl",
    "naniscript",
    "nasm",
    "neon",
    "nevod",
    "nginx",
    "nim",
    "nix",
    "nsis",
    "objectivec",
    "ocaml",
    "opencl",
    "openqasm",
    "oz",
    "parigp",
    "parser",
    "pascal",
    "pascaligo",
    "pcaxis",
    "peoplecode",
    "perl",
    "php",
    "php-extras",
    "phpdoc",
    "plsql",
    "powerquery",
    "powershell",
    "processing",
    "prolog",
    "promql",
    "properties",
    "protobuf",
    "psl",
    "pug",
    "puppet",
    "pure",
    "purebasic",
    "purescript",
    "python",
    "q",
    "qml",
    "qore",
    "qsharp",
    "r",
    "racket",
    "reason",
    "regex",
    "rego",
    "renpy",
    "rest",
    "rip",
    "roboconf",
    "robotframework",
    "ruby",
    "rust",
    "sas",
    "sass",
    "scala",
    "scheme",
    "scss",
    "shell-session",
    "smali",
    "smalltalk",
    "smarty",
    "sml",
    "solidity",
    "solution-file",
    "soy",
    "sparql",
    "splunk-spl",
    "sqf",
    "sql",
    "squirrel",
    "stan",
    "stylus",
    "swift",
    "systemd",
    "t4-cs",
    "t4-templating",
    "t4-vb",
    "tap",
    "tcl",
    "textile",
    "toml",
    "tremor",
    "tsx",
    "tt2",
    "turtle",
    "twig",
    "typescript",
    "typoscript",
    "unrealscript",
    "uorazor",
    "uri",
    "v",
    "vala",
    "vbnet",
    "velocity",
    "verilog",
    "vhdl",
    "vim",
    "visual-basic",
    "warpscript",
    "wasm",
    "web-idl",
    "wiki",
    "wolfram",
    "wren",
    "xeora",
    "xml-doc",
    "xojo",
    "xquery",
    "yaml",
    "yang",
    "zig",
]


def construct_theme_var(theme: str) -> Var[Theme]:
    """Construct a theme var.

    Args:
        theme: The theme to construct.

    Returns:
        The constructed theme var.
    """
    return Var(
        theme,
        _var_data=VarData(
            imports={
                f"react-syntax-highlighter/dist/esm/styles/prism/{format.to_kebab_case(theme)}": [
                    ImportVar(tag=theme, is_default=True, install=False)
                ]
            }
        ),
    )


@dataclasses.dataclass(init=False)
class Theme:
    """Themes for the CodeBlock component."""

    a11y_dark: ClassVar[Var[Theme]] = construct_theme_var("a11yDark")
    atom_dark: ClassVar[Var[Theme]] = construct_theme_var("atomDark")
    cb: ClassVar[Var[Theme]] = construct_theme_var("cb")
    coldark_cold: ClassVar[Var[Theme]] = construct_theme_var("coldarkCold")
    coldark_dark: ClassVar[Var[Theme]] = construct_theme_var("coldarkDark")
    coy: ClassVar[Var[Theme]] = construct_theme_var("coy")
    coy_without_shadows: ClassVar[Var[Theme]] = construct_theme_var("coyWithoutShadows")
    darcula: ClassVar[Var[Theme]] = construct_theme_var("darcula")
    dark: ClassVar[Var[Theme]] = construct_theme_var("oneDark")
    dracula: ClassVar[Var[Theme]] = construct_theme_var("dracula")
    duotone_dark: ClassVar[Var[Theme]] = construct_theme_var("duotoneDark")
    duotone_earth: ClassVar[Var[Theme]] = construct_theme_var("duotoneEarth")
    duotone_forest: ClassVar[Var[Theme]] = construct_theme_var("duotoneForest")
    duotone_light: ClassVar[Var[Theme]] = construct_theme_var("duotoneLight")
    duotone_sea: ClassVar[Var[Theme]] = construct_theme_var("duotoneSea")
    duotone_space: ClassVar[Var[Theme]] = construct_theme_var("duotoneSpace")
    funky: ClassVar[Var[Theme]] = construct_theme_var("funky")
    ghcolors: ClassVar[Var[Theme]] = construct_theme_var("ghcolors")
    gruvbox_dark: ClassVar[Var[Theme]] = construct_theme_var("gruvboxDark")
    gruvbox_light: ClassVar[Var[Theme]] = construct_theme_var("gruvboxLight")
    holi_theme: ClassVar[Var[Theme]] = construct_theme_var("holiTheme")
    hopscotch: ClassVar[Var[Theme]] = construct_theme_var("hopscotch")
    light: ClassVar[Var[Theme]] = construct_theme_var("oneLight")
    lucario: ClassVar[Var[Theme]] = construct_theme_var("lucario")
    material_dark: ClassVar[Var[Theme]] = construct_theme_var("materialDark")
    material_light: ClassVar[Var[Theme]] = construct_theme_var("materialLight")
    material_oceanic: ClassVar[Var[Theme]] = construct_theme_var("materialOceanic")
    night_owl: ClassVar[Var[Theme]] = construct_theme_var("nightOwl")
    nord: ClassVar[Var[Theme]] = construct_theme_var("nord")
    okaidia: ClassVar[Var[Theme]] = construct_theme_var("okaidia")
    one_dark: ClassVar[Var[Theme]] = construct_theme_var("oneDark")
    one_light: ClassVar[Var[Theme]] = construct_theme_var("oneLight")
    pojoaque: ClassVar[Var[Theme]] = construct_theme_var("pojoaque")
    prism: ClassVar[Var[Theme]] = construct_theme_var("prism")
    shades_of_purple: ClassVar[Var[Theme]] = construct_theme_var("shadesOfPurple")
    solarized_dark_atom: ClassVar[Var[Theme]] = construct_theme_var("solarizedDarkAtom")
    solarizedlight: ClassVar[Var[Theme]] = construct_theme_var("solarizedlight")
    synthwave84: ClassVar[Var[Theme]] = construct_theme_var("synthwave84")
    tomorrow: ClassVar[Var[Theme]] = construct_theme_var("tomorrow")
    twilight: ClassVar[Var[Theme]] = construct_theme_var("twilight")
    vs: ClassVar[Var[Theme]] = construct_theme_var("vs")
    vs_dark: ClassVar[Var[Theme]] = construct_theme_var("vsDark")
    vsc_dark_plus: ClassVar[Var[Theme]] = construct_theme_var("vscDarkPlus")
    xonokai: ClassVar[Var[Theme]] = construct_theme_var("xonokai")
    z_touch: ClassVar[Var[Theme]] = construct_theme_var("zTouch")


for theme_name in dir(Theme):
    if theme_name.startswith("_"):
        continue
    setattr(Theme, theme_name, getattr(Theme, theme_name)._replace(_var_type=Theme))


class CodeBlock(Component, MarkdownComponentMap):
    """A code block."""

    library = "react-syntax-highlighter@16.1.2"

    tag = "PrismAsyncLight"

    alias = "SyntaxHighlighter"

    theme: Var[Theme] = field(
        default=Theme.one_light,
        doc="The theme to use, accessed via rx.code_block.themes (e.g. rx.code_block.themes.one_dark).",
    )

    language: Var[LiteralCodeLanguage] = field(
        default=Var.create("python"), doc="The language to use."
    )

    code: Var[str] = field(doc="The code to display.")

    show_line_numbers: Var[bool] = field(
        doc="If this is enabled line numbers will be shown next to the code block."
    )

    starting_line_number: Var[int] = field(doc="The starting line number to use.")

    wrap_long_lines: Var[bool] = field(doc="Whether to wrap long lines.")

    custom_style: dict[str, str | Var | Color] = field(
        doc="A custom style for the code block.",
        default_factory=dict,
    )

    code_tag_props: Var[dict[str, str | dict[str, str]]] = field(
        doc="Props passed down to the code tag."
    )

    # Whether a copy button should appear.
    can_copy: bool | None = field(
        default=False,
        is_javascript_property=False,
    )

    # A custom copy button to override the default one.
    copy_button: bool | Component | None = field(
        default=None,
        is_javascript_property=False,
    )

    @classmethod
    def create(
        cls,
        *children,
        **props,
    ):
        """Create a text component.

        Args:
            *children: The children of the component.
            **props: The props to pass to the component.

        Returns:
            The text component.
        """
        from reflex_components_lucide.icon import Icon

        # This component handles style in a special prop.
        custom_style = props.pop("custom_style", {})
        can_copy = props.pop("can_copy", False)
        copy_button = props.pop("copy_button", None)

        # react-syntax-highlighter doesn't have an explicit "light" or "dark" theme so we use one-light and one-dark
        # themes respectively to ensure code compatibility.
        if "theme" not in props:
            # Default color scheme responds to global color mode.
            props["theme"] = color_mode_cond(
                light=Theme.one_light,
                dark=Theme.one_dark,
            )

        if props.get("wrap_long_lines") is True:
            code_tag_props = props.get("code_tag_props")
            if code_tag_props is None:
                props["code_tag_props"] = {"style": {"whiteSpace": "pre-wrap"}}
            elif isinstance(code_tag_props, dict):
                code_tag_props = code_tag_props.copy()
                code_tag_style = code_tag_props.get("style")
                if code_tag_style is None:
                    code_tag_props["style"] = {"whiteSpace": "pre-wrap"}
                elif (
                    isinstance(code_tag_style, dict)
                    and "whiteSpace" not in code_tag_style
                ):
                    code_tag_props["style"] = {
                        "whiteSpace": "pre-wrap",
                        **code_tag_style,
                    }
                props["code_tag_props"] = code_tag_props

        if can_copy:
            code = children[0]
            copy_button = (
                copy_button
                if copy_button is not None
                else Button.create(
                    Icon.create(tag="copy"),
                    aria_label="Copy code",
                    type="button",
                    on_click=set_clipboard(code),
                    style=Style({"position": "absolute", "top": "0.5em", "right": "0"}),
                )
            )
            custom_style.update({"padding": "1em 3.2em 1em 1em"})
        else:
            copy_button = None

        # Transfer style props to the custom style prop.
        for key, value in props.items():
            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"] = LiteralVar.create(props["code"])

        # Create the component.
        code_block = super().create(
            **props,
            custom_style=Style(custom_style),
        )

        if copy_button:
            return Box.create(code_block, copy_button, position="relative")
        return code_block

    def add_style(self):
        """Add style to the component."""
        self.custom_style.update(self.style)

    def _get_vars(
        self, include_children: bool = False, ignore_ids: set[int] | None = None
    ) -> Iterator[Var]:
        """Walk all Vars used in this component.

        Args:
            include_children: Whether to include Vars from children.
            ignore_ids: The ids to ignore.

        Yields:
            Each var referenced by the component, including any embedded in
            the custom_style dict, which is not a Var-typed prop.
        """
        yield from super()._get_vars(
            include_children=include_children, ignore_ids=ignore_ids
        )
        custom_style = self.custom_style
        if isinstance(custom_style, Style) and custom_style._var_data:
            yield Var(
                _js_expr="custom_style",
                _var_type=str,
                _var_data=custom_style._var_data,
            )

    def _render(self):
        out = super()._render()

        theme = self.theme

        return (
            out
            .add_props(style=theme)
            .remove_props("theme", "code")
            .add_props(
                children=self.code,
            )
        )


class CodeblockNamespace(ComponentNamespace):
    """Namespace for the CodeBlock component."""

    themes = Theme

    __call__ = CodeBlock.create


code_block = CodeblockNamespace()