Update CodeBlock class to accept rx.color in custom_style (#3168)

reflex/components/datadisplay/code.py

@@ -7,6 +7,7 @@ from reflex.components.chakra.layout import Box
 from reflex.components.chakra.media import Icon
 from reflex.components.component import Component
 from reflex.components.core.cond import color_mode_cond
+from reflex.constants.colors import Color
 from reflex.event import set_clipboard
 from reflex.style import Style
 from reflex.utils import format, imports
@@ -373,7 +374,7 @@ class CodeBlock(Component):
     wrap_long_lines: Var[bool]
 
     # A custom style for the code block.
-    custom_style: Dict[str, str] = {}
+    custom_style: Dict[str, Union[str, Var, Color]] = {}
 
     # Props passed down to the code tag.
     code_tag_props: Var[Dict[str, str]]

reflex/components/datadisplay/code.pyi

@@ -14,6 +14,7 @@ from reflex.components.chakra.layout import Box
 from reflex.components.chakra.media import Icon
 from reflex.components.component import Component
 from reflex.components.core.cond import color_mode_cond
+from reflex.constants.colors import Color
 from reflex.event import set_clipboard
 from reflex.style import Style
 from reflex.utils import format, imports
@@ -1029,7 +1030,7 @@ class CodeBlock(Component):
         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,
-        custom_style: Optional[Dict[str, str]] = None,
+        custom_style: Optional[Dict[str, Union[str, Var, Color]]] = None,
         code_tag_props: Optional[Union[Var[Dict[str, str]], Dict[str, str]]] = None,
         style: Optional[Style] = None,
         key: Optional[Any] = None,

tests/components/core/test_colors.py

@@ -1,6 +1,7 @@
 import pytest
 
 import reflex as rx
+from reflex.components.datadisplay.code import CodeBlock
 from reflex.vars import Var
 
 
@@ -90,3 +91,26 @@ def test_color(color, expected):
 )
 def test_color_with_conditionals(cond_var, expected):
     assert str(cond_var) == expected
+
+
+@pytest.mark.parametrize(
+    "color, expected",
+    [
+        (create_color_var(rx.color("red")), "var(--red-7)"),
+        (create_color_var(rx.color("green", shade=1)), "var(--green-1)"),
+        (create_color_var(rx.color("blue", alpha=True)), "var(--blue-a7)"),
+        ("red", "red"),
+        ("green", "green"),
+        ("blue", "blue"),
+    ],
+)
+def test_radix_color(color, expected):
+    """Test that custom_style can accept both string
+    literals and rx.color inputs.
+
+    Args:
+        color (Color): A Color made with rx.color
+        expected (str): The expected custom_style string, radix or literal
+    """
+    code_block = CodeBlock.create("Hello World", background_color=color)
+    assert code_block.custom_style["backgroundColor"].__format__("") == expected  # type: ignore