From 90b6a6657e43159f57cb9cf4fb636015ece99a16 Mon Sep 17 00:00:00 2001
From: Elijah <elijahahianyo@gmail.com>
Date: Fri, 1 Nov 2024 16:30:00 +0000
Subject: [PATCH] This should fix it. Right?

---
 .../components/datadisplay/shiki_code_block.py | 18 +++++++++++++++++-
 .../datadisplay/shiki_code_block.pyi           |  5 ++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/reflex/components/datadisplay/shiki_code_block.py b/reflex/components/datadisplay/shiki_code_block.py
index 4a3e05d0e..4c41dc743 100644
--- a/reflex/components/datadisplay/shiki_code_block.py
+++ b/reflex/components/datadisplay/shiki_code_block.py
@@ -12,6 +12,10 @@ from reflex.components.core.colors import color
 from reflex.components.core.cond import color_mode_cond
 from reflex.components.el.elements.forms import Button
 from reflex.components.lucide.icon import Icon
+from reflex.components.markdown.markdown import (
+    _LANGUAGE,
+    MarkdownComponentMap,
+)
 from reflex.components.props import NoExtrasAllowedProps
 from reflex.components.radix.themes.layout.box import Box
 from reflex.event import run_script, set_clipboard
@@ -528,7 +532,7 @@ class ShikiJsTransformer(ShikiBaseTransformers):
         super().__init__(**kwargs)
 
 
-class ShikiCodeBlock(Component):
+class ShikiCodeBlock(Component, MarkdownComponentMap):
     """A Code block."""
 
     library = "/components/shiki/code"
@@ -697,6 +701,18 @@ class ShikiCodeBlock(Component):
 
         return processed
 
+    @classmethod
+    def get_component_map_custom_code(cls) -> str:
+        """Get the custom code for the component.
+
+        Returns:
+            The custom code for the component.
+        """
+        return f"""
+    const match = (className || '').match(/language-(?<lang>.*)/);
+    const {str(_LANGUAGE)} = match ? match[1] : '';
+    """
+
 
 class ShikiHighLevelCodeBlock(ShikiCodeBlock):
     """High level component for the shiki syntax highlighter."""
diff --git a/reflex/components/datadisplay/shiki_code_block.pyi b/reflex/components/datadisplay/shiki_code_block.pyi
index 543030daa..f75cefd4d 100644
--- a/reflex/components/datadisplay/shiki_code_block.pyi
+++ b/reflex/components/datadisplay/shiki_code_block.pyi
@@ -7,6 +7,7 @@ from typing import Any, Dict, Literal, Optional, Union, overload
 
 from reflex.base import Base
 from reflex.components.component import Component, ComponentNamespace
+from reflex.components.markdown.markdown import MarkdownComponentMap
 from reflex.components.props import NoExtrasAllowedProps
 from reflex.event import BASE_STATE, EventType
 from reflex.style import Style
@@ -350,7 +351,7 @@ class ShikiJsTransformer(ShikiBaseTransformers):
     fns: list[FunctionStringVar]
     style: Optional[Style]
 
-class ShikiCodeBlock(Component):
+class ShikiCodeBlock(Component, MarkdownComponentMap):
     @overload
     @classmethod
     def create(  # type: ignore
@@ -972,6 +973,8 @@ class ShikiCodeBlock(Component):
     def create_transformer(
         cls, library: str, fns: list[str]
     ) -> ShikiBaseTransformers: ...
+    @classmethod
+    def get_component_map_custom_code(cls) -> str: ...
 
 class ShikiHighLevelCodeBlock(ShikiCodeBlock):
     @overload