shiki_code_block.py

"""Shiki syntax hghlighter component."""

from __future__ import annotations

import dataclasses
import re
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Literal

from reflex_base.components.component import Component, ComponentNamespace, field
from reflex_base.components.props import NoExtrasAllowedProps
from reflex_base.event import run_script, set_clipboard
from reflex_base.style import Style
from reflex_base.utils.exceptions import VarTypeError
from reflex_base.utils.imports import ImportVar
from reflex_base.vars.base import LiteralVar, Var
from reflex_base.vars.function import FunctionStringVar
from reflex_base.vars.sequence import StringVar, string_replace_operation
from reflex_components_core.core.colors import color
from reflex_components_core.core.cond import color_mode_cond
from reflex_components_core.core.markdown_component_map import MarkdownComponentMap
from reflex_components_core.el.elements.forms import Button
from reflex_components_radix.themes.layout.box import Box


def copy_script() -> Any:
    """Copy script for the code block and modify the child SVG element.

    Returns:
        Any: The result of calling the script.
    """
    return run_script(
        """
// Event listener for the parent click
document.addEventListener('click', function(event) {
    // Find the closest button (parent element)
    const parent = event.target.closest('button');
    // If the parent is found
    if (parent) {
        // Find the SVG element within the parent
        const svgIcon = parent.querySelector('svg');
        // If the SVG exists, proceed with the script
        if (svgIcon) {
            const originalPath = svgIcon.innerHTML;
            const checkmarkPath = '<polyline points="20 6 9 17 4 12"></polyline>';  // Checkmark SVG path
            function transition(element, scale, opacity) {
                element.style.transform = `scale(${scale})`;
                element.style.opacity = opacity;
            }
            // Animate the SVG
            transition(svgIcon, 0, '0');
            setTimeout(() => {
                svgIcon.innerHTML = checkmarkPath;  // Replace content with checkmark
                svgIcon.setAttribute('viewBox', '0 0 24 24');  // Adjust viewBox if necessary
                transition(svgIcon, 1, '1');
                setTimeout(() => {
                    transition(svgIcon, 0, '0');
                    setTimeout(() => {
                        svgIcon.innerHTML = originalPath;  // Restore original SVG content
                        transition(svgIcon, 1, '1');
                    }, 60);
                }, 350);
            }, 60);
        } else {
            // console.error('SVG element not found within the parent.');
        }
    } else {
        // console.error('Parent element not found.');
    }
})
"""
    )


SHIKIJS_TRANSFORMER_FNS = {
    "transformerNotationDiff",
    "transformerNotationHighlight",
    "transformerNotationWordHighlight",
    "transformerNotationFocus",
    "transformerNotationErrorLevel",
    "transformerRenderWhitespace",
    "transformerMetaHighlight",
    "transformerMetaWordHighlight",
    "transformerCompactLineOptions",
    # TODO: this transformer when included adds a weird behavior which removes other code lines. Need to figure out why.
    # "transformerRemoveLineBreak",
    "transformerRemoveNotationEscape",
}
LINE_NUMBER_STYLING = {
    "code": {
        "counter-reset": "step",
        "counter-increment": "step 0",
        "display": "grid",
        "line-height": "1.7",
        "font-size": "0.875em",
    },
    "code .line::before": {
        "content": "counter(step)",
        "counter-increment": "step",
        "width": "1rem",
        "margin-right": "1.5rem",
        "display": "inline-block",
        "text-align": "right",
        "color": "rgba(115,138,148,.4)",
    },
}
BOX_PARENT_STYLING = {
    "pre": {
        "margin": "0",
        "padding": "24px",
        "background": "transparent",
        "overflow-x": "auto",
        "border-radius": "6px",
    },
}

THEME_MAPPING = {
    "light": "one-light",
    "dark": "one-dark-pro",
    "a11y-dark": "github-dark",
}
LANGUAGE_MAPPING = {"bash": "shellscript"}
LiteralCodeLanguage = Literal[
    "abap",
    "actionscript-3",
    "ada",
    "angular-html",
    "angular-ts",
    "apache",
    "apex",
    "apl",
    "applescript",
    "ara",
    "asciidoc",
    "asm",
    "astro",
    "awk",
    "ballerina",
    "bat",
    "beancount",
    "berry",
    "bibtex",
    "bicep",
    "blade",
    "c",
    "cadence",
    "clarity",
    "clojure",
    "cmake",
    "cobol",
    "codeowners",
    "codeql",
    "coffee",
    "common-lisp",
    "coq",
    "cpp",
    "crystal",
    "csharp",
    "css",
    "csv",
    "cue",
    "cypher",
    "d",
    "dart",
    "dax",
    "desktop",
    "diff",
    "docker",
    "dotenv",
    "dream-maker",
    "edge",
    "elixir",
    "elm",
    "emacs-lisp",
    "erb",
    "erlang",
    "fennel",
    "fish",
    "fluent",
    "fortran-fixed-form",
    "fortran-free-form",
    "fsharp",
    "gdresource",
    "gdscript",
    "gdshader",
    "genie",
    "gherkin",
    "git-commit",
    "git-rebase",
    "gleam",
    "glimmer-js",
    "glimmer-ts",
    "glsl",
    "gnuplot",
    "go",
    "graphql",
    "groovy",
    "hack",
    "haml",
    "handlebars",
    "haskell",
    "haxe",
    "hcl",
    "hjson",
    "hlsl",
    "html",
    "html-derivative",
    "http",
    "hxml",
    "hy",
    "imba",
    "ini",
    "java",
    "javascript",
    "jinja",
    "jison",
    "json",
    "json5",
    "jsonc",
    "jsonl",
    "jsonnet",
    "jssm",
    "jsx",
    "julia",
    "kotlin",
    "kusto",
    "latex",
    "lean",
    "less",
    "liquid",
    "log",
    "logo",
    "lua",
    "luau",
    "make",
    "markdown",
    "marko",
    "matlab",
    "mdc",
    "mdx",
    "mermaid",
    "mojo",
    "move",
    "narrat",
    "nextflow",
    "nginx",
    "nim",
    "nix",
    "nushell",
    "objective-c",
    "objective-cpp",
    "ocaml",
    "pascal",
    "perl",
    "php",
    "plain",
    "plsql",
    "po",
    "postcss",
    "powerquery",
    "powershell",
    "prisma",
    "prolog",
    "proto",
    "pug",
    "puppet",
    "purescript",
    "python",
    "qml",
    "qmldir",
    "qss",
    "r",
    "racket",
    "raku",
    "razor",
    "reg",
    "regexp",
    "rel",
    "riscv",
    "rst",
    "ruby",
    "rust",
    "sas",
    "sass",
    "scala",
    "scheme",
    "scss",
    "shaderlab",
    "shellscript",
    "shellsession",
    "smalltalk",
    "solidity",
    "soy",
    "sparql",
    "splunk",
    "sql",
    "ssh-config",
    "stata",
    "stylus",
    "svelte",
    "swift",
    "system-verilog",
    "systemd",
    "tasl",
    "tcl",
    "templ",
    "terraform",
    "tex",
    "toml",
    "ts-tags",
    "tsv",
    "tsx",
    "turtle",
    "twig",
    "typescript",
    "typespec",
    "typst",
    "v",
    "vala",
    "vb",
    "verilog",
    "vhdl",
    "viml",
    "vue",
    "vue-html",
    "vyper",
    "wasm",
    "wenyan",
    "wgsl",
    "wikitext",
    "wolfram",
    "xml",
    "xsl",
    "yaml",
    "zenscript",
    "zig",
]
LiteralCodeTheme = Literal[
    "andromeeda",
    "aurora-x",
    "ayu-dark",
    "catppuccin-frappe",
    "catppuccin-latte",
    "catppuccin-macchiato",
    "catppuccin-mocha",
    "dark-plus",
    "dracula",
    "dracula-soft",
    "everforest-dark",
    "everforest-light",
    "github-dark",
    "github-dark-default",
    "github-dark-dimmed",
    "github-dark-high-contrast",
    "github-light",
    "github-light-default",
    "github-light-high-contrast",
    "houston",
    "laserwave",
    "light-plus",
    "material-theme",
    "material-theme-darker",
    "material-theme-lighter",
    "material-theme-ocean",
    "material-theme-palenight",
    "min-dark",
    "min-light",
    "monokai",
    "night-owl",
    "nord",
    "one-dark-pro",
    "one-light",
    "plastic",
    "poimandres",
    "red",
    # rose-pine themes dont work with the current version of shikijs transformers
    # https://github.com/shikijs/shiki/issues/730
    "rose-pine",
    "rose-pine-dawn",
    "rose-pine-moon",
    "slack-dark",
    "slack-ochin",
    "snazzy-light",
    "solarized-dark",
    "solarized-light",
    "synthwave-84",
    "tokyo-night",
    "vesper",
    "vitesse-black",
    "vitesse-dark",
    "vitesse-light",
]


class Position(NoExtrasAllowedProps):
    """Position of the decoration."""

    line: int
    character: int


class ShikiDecorations(NoExtrasAllowedProps):
    """Decorations for the code block."""

    start: int | Position
    end: int | Position
    tag_name: str = "span"
    properties: dict[str, Any] = {}
    always_wrap: bool = False


@dataclass(kw_only=True)
class ShikiBaseTransformers:
    """Base for creating transformers."""

    library: str = ""
    fns: list[FunctionStringVar] = dataclasses.field(default_factory=list)
    style: Style | None = dataclasses.field(default=None)


@dataclass(kw_only=True)
class ShikiJsTransformer(ShikiBaseTransformers):
    """A Wrapped shikijs transformer."""

    library: str = "@shikijs/transformers@3.3.0"
    fns: list[FunctionStringVar] = dataclasses.field(
        default_factory=lambda: [
            FunctionStringVar.create(fn) for fn in SHIKIJS_TRANSFORMER_FNS
        ]
    )
    style: Style | None = dataclasses.field(
        default_factory=lambda: Style({
            "code": {
                "line-height": "1.7",
                "font-size": "0.875em",
                "display": "grid",
            },
            # Diffs
            ".diff": {
                "margin": "0 -24px",
                "padding": "0 24px",
                "width": "calc(100% + 48px)",
                "display": "inline-block",
            },
            ".diff.add": {
                "background-color": "rgba(16, 185, 129, .14)",
                "position": "relative",
            },
            ".diff.remove": {
                "background-color": "rgba(244, 63, 94, .14)",
                "opacity": "0.7",
                "position": "relative",
            },
            ".diff.remove:after": {
                "position": "absolute",
                "left": "10px",
                "content": "'-'",
                "color": "#b34e52",
            },
            ".diff.add:after": {
                "position": "absolute",
                "left": "10px",
                "content": "'+'",
                "color": "#18794e",
            },
            # Highlight
            ".highlighted": {
                "background-color": "rgba(142, 150, 170, .14)",
                "margin": "0 -24px",
                "padding": "0 24px",
                "width": "calc(100% + 48px)",
                "display": "inline-block",
            },
            ".highlighted.error": {
                "background-color": "rgba(244, 63, 94, .14)",
            },
            ".highlighted.warning": {
                "background-color": "rgba(234, 179, 8, .14)",
            },
            # Highlighted Word
            ".highlighted-word": {
                "background-color": color("gray", 2),
                "border": f"1px solid {color('gray', 5)}",
                "padding": "1px 3px",
                "margin": "-1px -3px",
                "border-radius": "4px",
            },
            # Focused Lines
            ".has-focused .line:not(.focused)": {
                "opacity": "0.7",
                "filter": "blur(0.095rem)",
                "transition": "filter .35s, opacity .35s",
            },
            ".has-focused:hover .line:not(.focused)": {
                "opacity": "1",
                "filter": "none",
            },
            # White Space
            # ".tab, .space": {
            #     "position": "relative", # noqa: ERA001
            # },
            # ".tab::before": {
            #     "content": "'⇥'", # noqa: ERA001
            #     "position": "absolute", # noqa: ERA001
            #     "opacity": "0.3",# noqa: ERA001
            # },
            # ".space::before": {
            #     "content": "'·'", # noqa: ERA001
            #     "position": "absolute", # noqa: ERA001
            #     "opacity": "0.3", # noqa: ERA001
            # },
        })
    )

    def __init__(self, **kwargs):
        """Initialize the transformer.

        Args:
            kwargs: Kwargs to initialize the props.

        """
        fns = kwargs.pop("fns", None)
        style = kwargs.pop("style", None)
        if fns:
            kwargs["fns"] = [
                (
                    FunctionStringVar.create(x)
                    if not isinstance(x, FunctionStringVar)
                    else x
                )
                for x in fns
            ]

        if style:
            kwargs["style"] = Style(style)
        super().__init__(**kwargs)


class ShikiCodeBlock(Component, MarkdownComponentMap):
    """A Code block."""

    library = "/components/shiki/code"

    tag = "Code"

    alias = "ShikiCode"

    lib_dependencies: list[str] = ["shiki@3.3.0"]

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

    theme: Var[LiteralCodeTheme] = field(
        default=Var.create("one-light"), doc='The theme to use ("light" or "dark").'
    )

    themes: Var[list[dict[str, Any]] | dict[str, str]] = field(
        doc="The set of themes to use for different modes."
    )

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

    transformers: Var[list[ShikiBaseTransformers | dict[str, Any]]] = field(
        default=Var.create([]),
        doc="The transformers to use for the syntax highlighter.",
    )

    decorations: Var[list[ShikiDecorations]] = field(
        default=Var.create([]), doc="The decorations to use for the syntax highlighter."
    )

    @classmethod
    def create(
        cls,
        *children,
        **props,
    ) -> Component:
        """Create a code block component using [shiki syntax highlighter](https://shiki.matsu.io/).

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

        Returns:
            The code block component.
        """
        # Separate props for the code block and the wrapper
        code_block_props = {}
        code_wrapper_props = {}
        decorations = props.pop("decorations", [])

        class_props = cls.get_props()

        # Distribute props between the code block and wrapper
        for key, value in props.items():
            (code_block_props if key in class_props else code_wrapper_props)[key] = (
                value
            )

        # cast decorations into ShikiDecorations.
        decorations = [
            ShikiDecorations(**decoration)
            if not isinstance(decoration, ShikiDecorations)
            else decoration
            for decoration in decorations
        ]
        code_block_props["decorations"] = decorations

        code_block_props["code"] = children[0]
        code_block = super().create(**code_block_props)

        transformer_styles = {}
        # Collect styles from transformers and wrapper
        for transformer in code_block.transformers._var_value:  # pyright: ignore [reportAttributeAccessIssue]
            if isinstance(transformer, ShikiBaseTransformers) and transformer.style:
                transformer_styles.update(transformer.style)
        transformer_styles.update(code_wrapper_props.pop("style", {}))

        return Box.create(
            code_block,
            *children[1:],
            style=Style({**transformer_styles, **BOX_PARENT_STYLING}),
            **code_wrapper_props,
        )

    def add_imports(self) -> dict[str, list[str]]:
        """Add the necessary imports.
        We add all referenced transformer functions as imports from their corresponding
        libraries.

        Returns:
            Imports for the component.

        Raises:
            ValueError: If the transformers are not of type LiteralVar.
        """
        imports = defaultdict(list)
        if not isinstance(self.transformers, LiteralVar):
            msg = f"transformers should be a LiteralVar type. Got {type(self.transformers)} instead."
            raise ValueError(msg)
        for transformer in self.transformers._var_value:
            if isinstance(transformer, ShikiBaseTransformers):
                imports[transformer.library].extend([
                    ImportVar(tag=str(fn)) for fn in transformer.fns
                ])
                if transformer.library not in self.lib_dependencies:
                    self.lib_dependencies.append(transformer.library)
        return imports

    @classmethod
    def create_transformer(cls, library: str, fns: list[str]) -> ShikiBaseTransformers:
        """Create a transformer from a third party library.

        Args:
            library: The name of the library.
            fns: The str names of the functions/callables to invoke from the library.

        Returns:
            A transformer for the specified library.

        Raises:
            ValueError: If a supplied function name is not valid str.
        """
        if any(not isinstance(fn_name, str) for fn_name in fns):
            msg = f"the function names should be str names of functions in the specified transformer: {library!r}"
            raise ValueError(msg)
        return ShikiBaseTransformers(
            library=library,
            fns=[FunctionStringVar.create(fn) for fn in fns],  # pyright: ignore [reportCallIssue]
        )

    def _render(self, props: dict[str, Any] | None = None):
        """Renders the component with the given properties, processing transformers if present.

        Args:
            props: Optional properties to pass to the render function.

        Returns:
            Rendered component output.
        """
        # Ensure props is initialized from class attributes if not provided
        props = props or {
            attr.rstrip("_"): getattr(self, attr) for attr in self.get_props()
        }

        # Extract transformers and apply transformations
        transformers = props.get("transformers")
        if transformers is not None:
            transformed_values = self._process_transformers(transformers._var_value)
            props["transformers"] = LiteralVar.create(transformed_values)

        return super()._render(props)

    def _process_transformers(self, transformer_list: list) -> list:
        """Processes a list of transformers, applying transformations where necessary.

        Args:
            transformer_list: List of transformer objects or values.

        Returns:
            list: A list of transformed values.
        """
        processed = []

        for transformer in transformer_list:
            if isinstance(transformer, ShikiBaseTransformers):
                processed.extend(fn.call() for fn in transformer.fns)
            else:
                processed.append(transformer)

        return processed


class ShikiHighLevelCodeBlock(ShikiCodeBlock):
    """High level component for the shiki syntax highlighter."""

    use_transformers: Var[bool] = field(
        doc="If this is enabled, the default transformers(shikijs transformer) will be used."
    )

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

    can_copy: bool = field(
        doc="Whether a copy button should appear.",
        default=False,
        is_javascript_property=False,
    )

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

    @classmethod
    def create(
        cls,
        *children,
        **props,
    ) -> Component:
        """Create a code block component using [shiki syntax highlighter](https://shiki.matsu.io/).

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

        Returns:
            The code block component.
        """
        from reflex_components_lucide.icon import Icon

        use_transformers = props.pop("use_transformers", False)
        show_line_numbers = props.pop("show_line_numbers", False)
        language = props.pop("language", None)
        can_copy = props.pop("can_copy", False)
        copy_button = props.pop("copy_button", None)

        if use_transformers:
            props["transformers"] = [ShikiJsTransformer()]

        if language is not None:
            props["language"] = cls._map_languages(language)

        # line numbers are generated via css
        if show_line_numbers:
            props["style"] = {**LINE_NUMBER_STYLING, **props.get("style", {})}

        theme = props.pop("theme", None)
        props["theme"] = props["theme"] = (
            cls._map_themes(theme)
            if theme is not None
            else color_mode_cond(  # Default color scheme responds to global color mode.
                light="one-light",
                dark="one-dark-pro",
            )
        )

        if can_copy:
            code = children[0]
            copy_button = (
                copy_button
                if copy_button is not None
                else Button.create(
                    Icon.create(tag="copy", size=16, color=color("gray", 11)),
                    aria_label="Copy code",
                    on_click=[
                        set_clipboard(cls._strip_transformer_triggers(code)),
                        copy_script(),
                    ],
                    style=Style({
                        "position": "absolute",
                        "top": "4px",
                        "right": "4px",
                        "background": color("gray", 3),
                        "border": "1px solid",
                        "border-color": color("gray", 5),
                        "border-radius": "6px",
                        "padding": "5px",
                        "opacity": "1",
                        "cursor": "pointer",
                        "_hover": {
                            "background": color("gray", 4),
                        },
                        "transition": "background 0.12s ease-out",
                        "&>svg": {
                            "transition": "transform 0.12s ease-out, opacity 0.12s ease-out",
                        },
                        "_active": {
                            "background": color("gray", 5),
                        },
                    }),
                )
            )

        if copy_button:
            return ShikiCodeBlock.create(
                children[0], copy_button, position="relative", **props
            )
        return ShikiCodeBlock.create(children[0], **props)

    @staticmethod
    def _map_themes(theme: str) -> str:
        if isinstance(theme, str) and theme in THEME_MAPPING:
            return THEME_MAPPING[theme]
        return theme

    @staticmethod
    def _map_languages(language: str) -> str:
        if isinstance(language, str) and language in LANGUAGE_MAPPING:
            return LANGUAGE_MAPPING[language]
        return language

    @staticmethod
    def _strip_transformer_triggers(code: str | StringVar) -> StringVar | str:
        if not isinstance(code, (StringVar, str)):
            msg = f"code should be string literal or a StringVar type. Got {type(code)} instead."
            raise VarTypeError(msg)
        regex_pattern = r"[\/#]+ *\[!code.*?\]"

        if isinstance(code, Var):
            return string_replace_operation(
                code, StringVar(_js_expr=f"/{regex_pattern}/g", _var_type=str), ""
            )
        if isinstance(code, str):
            return re.sub(regex_pattern, "", code)
        return None


class TransformerNamespace(ComponentNamespace):
    """Namespace for the Transformers."""

    shikijs = ShikiJsTransformer


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

    root = staticmethod(ShikiCodeBlock.create)
    create_transformer = staticmethod(ShikiCodeBlock.create_transformer)
    transformers = TransformerNamespace()
    __call__ = staticmethod(ShikiHighLevelCodeBlock.create)


code_block = CodeblockNamespace()