add allow_system prop to colormode iconbutton, and clean up logic (#3507)

* add allow_system prop to colormode iconbutton, and clean up logic

* remove segmentedcontrol change from this PR

* make it work for chakraColorProvider too

* add comment to explain resolved_color_mode

reflex/.templates/web/components/reflex/chakra_color_mode_provider.js

@@ -1,21 +1,35 @@
-import { useColorMode as chakraUseColorMode } from "@chakra-ui/react"
-import { useTheme } from "next-themes"
-import { useEffect } from "react"
-import { ColorModeContext } from "/utils/context.js"
+import { useColorMode as chakraUseColorMode } from "@chakra-ui/react";
+import { useTheme } from "next-themes";
+import { useEffect, useState } from "react";
+import { ColorModeContext, defaultColorMode } from "/utils/context.js";
 
 export default function ChakraColorModeProvider({ children }) {
-  const {colorMode, toggleColorMode} = chakraUseColorMode()
-  const {theme, setTheme} = useTheme()
+  const { theme, resolvedTheme, setTheme } = useTheme();
+  const { colorMode, toggleColorMode } = chakraUseColorMode();
+  const [resolvedColorMode, setResolvedColorMode] = useState(theme);
 
   useEffect(() => {
-    if (colorMode != theme) {
-        toggleColorMode()
+    if (colorMode != resolvedTheme) {
+      toggleColorMode();
     }
-  }, [theme])
+  }, [theme, resolvedTheme]);
 
+  const rawColorMode = colorMode;
+  const setColorMode = (mode) => {
+    const allowedModes = ["light", "dark", "system"];
+    if (!allowedModes.includes(mode)) {
+      console.error(
+        `Invalid color mode "${mode}". Defaulting to "${defaultColorMode}".`
+      );
+      mode = defaultColorMode;
+    }
+    setTheme(mode);
+  };
   return (
-    <ColorModeContext.Provider value={[ colorMode, toggleColorMode ]}>
+    <ColorModeContext.Provider
+      value={{ rawColorMode, resolvedColorMode, toggleColorMode, setColorMode }}
+    >
       {children}
     </ColorModeContext.Provider>
-  )
-}
+  );
+}

reflex/.templates/web/components/reflex/radix_themes_color_mode_provider.js

@@ -3,18 +3,32 @@ import { useEffect, useState } from "react";
 import { ColorModeContext, defaultColorMode } from "/utils/context.js";
 
 export default function RadixThemesColorModeProvider({ children }) {
-  const { resolvedTheme, setTheme } = useTheme();
-  const [colorMode, setColorMode] = useState(defaultColorMode);
+  const { theme, resolvedTheme, setTheme } = useTheme();
+  const [rawColorMode, setRawColorMode] = useState(defaultColorMode);
+  const [resolvedColorMode, setResolvedColorMode] = useState(theme);
 
   useEffect(() => {
-    setColorMode(resolvedTheme);
-  }, [resolvedTheme]);
+    setRawColorMode(theme);
+    setResolvedColorMode(resolvedTheme);
+  }, [theme, resolvedTheme]);
 
   const toggleColorMode = () => {
     setTheme(resolvedTheme === "light" ? "dark" : "light");
   };
+  const setColorMode = (mode) => {
+    const allowedModes = ["light", "dark", "system"];
+    if (!allowedModes.includes(mode)) {
+      console.error(
+        `Invalid color mode "${mode}". Defaulting to "${defaultColorMode}".`
+      );
+      mode = defaultColorMode;
+    }
+    setTheme(mode);
+  };
   return (
-    <ColorModeContext.Provider value={[colorMode, toggleColorMode]}>
+    <ColorModeContext.Provider
+      value={{ rawColorMode, resolvedColorMode, toggleColorMode, setColorMode }}
+    >
       {children}
     </ColorModeContext.Provider>
   );

reflex/compiler/templates.py

@@ -37,7 +37,9 @@ class ReflexJinjaEnvironment(Environment):
                 constants.CompileVars.PROCESSING: False,
             },
             "color_mode": constants.ColorMode.NAME,
+            "resolved_color_mode": constants.ColorMode.RESOLVED_NAME,
             "toggle_color_mode": constants.ColorMode.TOGGLE,
+            "set_color_mode": constants.ColorMode.SET,
             "use_color_mode": constants.ColorMode.USE,
             "hydrate": constants.CompileVars.HYDRATE,
             "on_load_internal": constants.CompileVars.ON_LOAD_INTERNAL,

reflex/components/core/cond.py

@@ -9,7 +9,7 @@ from reflex.components.component import BaseComponent, Component, MemoizationLea
 from reflex.components.tags import CondTag, Tag
 from reflex.constants import Dirs
 from reflex.constants.colors import Color
-from reflex.style import LIGHT_COLOR_MODE, color_mode
+from reflex.style import LIGHT_COLOR_MODE, resolved_color_mode
 from reflex.utils import format
 from reflex.utils.imports import ImportDict, ImportVar
 from reflex.vars import Var, VarData
@@ -208,7 +208,7 @@ def color_mode_cond(light: Any, dark: Any = None) -> Var | Component:
         The conditional component or prop.
     """
     return cond(
-        color_mode == Var.create(LIGHT_COLOR_MODE, _var_is_string=True),
+        resolved_color_mode == Var.create(LIGHT_COLOR_MODE, _var_is_string=True),
         light,
         dark,
     )

reflex/components/radix/themes/color_mode.py

@@ -23,8 +23,10 @@ from typing import Literal, get_args
 from reflex.components.component import BaseComponent
 from reflex.components.core.cond import Cond, color_mode_cond, cond
 from reflex.components.lucide.icon import Icon
+from reflex.components.radix.themes.components.dropdown_menu import dropdown_menu
 from reflex.components.radix.themes.components.switch import Switch
-from reflex.style import LIGHT_COLOR_MODE, color_mode, toggle_color_mode
+from reflex.event import EventChain
+from reflex.style import LIGHT_COLOR_MODE, color_mode, set_color_mode, toggle_color_mode
 from reflex.utils import console
 from reflex.vars import BaseVar, Var
 
@@ -95,6 +97,7 @@ class ColorModeIconButton(IconButton):
         cls,
         *children,
         position: LiteralPosition | None = None,
+        allow_system: bool = False,
         **props,
     ):
         """Create a icon button component that calls toggle_color_mode on click.
@@ -102,6 +105,7 @@ class ColorModeIconButton(IconButton):
         Args:
             *children: The children of the component.
             position: The position of the icon button. Follow document flow if None.
+            allow_system: Allow picking the "system" value for the color mode.
             **props: The props to pass to the component.
 
         Returns:
@@ -137,6 +141,32 @@ class ColorModeIconButton(IconButton):
         props.setdefault("z_index", "20")
         props.setdefault(":hover", {"cursor": "pointer"})
 
+        if allow_system:
+
+            def color_mode_item(_color_mode):
+                setter = Var.create_safe(
+                    f'() => {set_color_mode._var_name}("{_color_mode}")',
+                    _var_is_string=False,
+                    _var_is_local=True,
+                    _var_data=set_color_mode._var_data,
+                )
+                setter._var_type = EventChain
+
+                return dropdown_menu.item(_color_mode.title(), on_click=setter)  # type: ignore
+
+            return dropdown_menu.root(
+                dropdown_menu.trigger(
+                    super().create(
+                        ColorModeIcon.create(),
+                        **props,
+                    )
+                ),
+                dropdown_menu.content(
+                    color_mode_item("light"),
+                    color_mode_item("dark"),
+                    color_mode_item("system"),
+                ),
+            )
         return super().create(
             ColorModeIcon.create(),
             on_click=toggle_color_mode,

reflex/components/radix/themes/color_mode.pyi

@@ -12,8 +12,10 @@ from typing import Literal, get_args
 from reflex.components.component import BaseComponent
 from reflex.components.core.cond import Cond, color_mode_cond, cond
 from reflex.components.lucide.icon import Icon
+from reflex.components.radix.themes.components.dropdown_menu import dropdown_menu
 from reflex.components.radix.themes.components.switch import Switch
-from reflex.style import LIGHT_COLOR_MODE, color_mode, toggle_color_mode
+from reflex.event import EventChain
+from reflex.style import LIGHT_COLOR_MODE, color_mode, set_color_mode, toggle_color_mode
 from reflex.utils import console
 from reflex.vars import BaseVar, Var
 from .components.icon_button import IconButton
@@ -113,6 +115,7 @@ class ColorModeIconButton(IconButton):
         position: Optional[
             Literal["top-left", "top-right", "bottom-left", "bottom-right"]
         ] = None,
+        allow_system: Optional[bool] = False,
         as_child: Optional[Union[Var[bool], bool]] = None,
         size: Optional[
             Union[Var[Literal["1", "2", "3", "4"]], Literal["1", "2", "3", "4"]]
@@ -316,6 +319,7 @@ class ColorModeIconButton(IconButton):
         Args:
             *children: The children of the component.
             position: The position of the icon button. Follow document flow if None.
+            allow_system: Allow picking the "system" value for the color mode.
             as_child: Change the default rendered element for the one passed as a child, merging their props and behavior.
             size: Button size "1" - "4"
             variant: Variant of button: "classic" | "solid" | "soft" | "surface" | "outline" | "ghost"

reflex/components/sonner/toast.py

@@ -12,7 +12,7 @@ from reflex.event import (
     EventSpec,
     call_script,
 )
-from reflex.style import Style, color_mode
+from reflex.style import Style, resolved_color_mode
 from reflex.utils import format
 from reflex.utils.imports import ImportVar
 from reflex.utils.serializers import serialize, serializer
@@ -168,7 +168,7 @@ class Toaster(Component):
     tag = "Toaster"
 
     # the theme of the toast
-    theme: Var[str] = color_mode
+    theme: Var[str] = resolved_color_mode
 
     # whether to show rich colors
     rich_colors: Var[bool] = Var.create_safe(True)

reflex/components/sonner/toast.pyi

@@ -13,7 +13,7 @@ from reflex.components.component import Component, ComponentNamespace
 from reflex.components.lucide.icon import Icon
 from reflex.components.props import PropsBase
 from reflex.event import EventSpec, call_script
-from reflex.style import Style, color_mode
+from reflex.style import Style, resolved_color_mode
 from reflex.utils import format
 from reflex.utils.imports import ImportVar
 from reflex.utils.serializers import serialize, serializer

reflex/constants/base.py

@@ -126,9 +126,11 @@ class Next(SimpleNamespace):
 class ColorMode(SimpleNamespace):
     """Constants related to ColorMode."""
 
-    NAME = "colorMode"
+    NAME = "rawColorMode"
+    RESOLVED_NAME = "resolvedColorMode"
     USE = "useColorMode"
     TOGGLE = "toggleColorMode"
+    SET = "setColorMode"
 
 
 # Env modes

reflex/style.py

@@ -17,27 +17,40 @@ LIGHT_COLOR_MODE: str = "light"
 DARK_COLOR_MODE: str = "dark"
 
 # Reference the global ColorModeContext
-color_mode_var_data = VarData(
-    imports={
-        f"/{constants.Dirs.CONTEXTS_PATH}": [ImportVar(tag="ColorModeContext")],
-        "react": [ImportVar(tag="useContext")],
-    },
-    hooks={
-        f"const [ {constants.ColorMode.NAME}, {constants.ColorMode.TOGGLE} ] = useContext(ColorModeContext)": None,
-    },
-)
-# Var resolves to the current color mode for the app ("light" or "dark")
+color_mode_imports = {
+    f"/{constants.Dirs.CONTEXTS_PATH}": [ImportVar(tag="ColorModeContext")],
+    "react": [ImportVar(tag="useContext")],
+}
+color_mode_toggle_hooks = {
+    f"const {{ {constants.ColorMode.RESOLVED_NAME}, {constants.ColorMode.TOGGLE} }} = useContext(ColorModeContext)": None,
+}
+color_mode_set_hooks = {
+    f"const {{ {constants.ColorMode.NAME}, {constants.ColorMode.RESOLVED_NAME}, {constants.ColorMode.TOGGLE}, {constants.ColorMode.SET} }} = useContext(ColorModeContext)": None,
+}
+color_mode_var_data = VarData(imports=color_mode_imports, hooks=color_mode_toggle_hooks)
+# Var resolves to the current color mode for the app ("light", "dark" or "system")
 color_mode = BaseVar(
     _var_name=constants.ColorMode.NAME,
     _var_type="str",
     _var_data=color_mode_var_data,
 )
+# Var resolves to the resolved color mode for the app ("light" or "dark")
+resolved_color_mode = BaseVar(
+    _var_name=constants.ColorMode.RESOLVED_NAME,
+    _var_type="str",
+    _var_data=color_mode_var_data,
+)
 # Var resolves to a function invocation that toggles the color mode
 toggle_color_mode = BaseVar(
     _var_name=constants.ColorMode.TOGGLE,
     _var_type=EventChain,
     _var_data=color_mode_var_data,
 )
+set_color_mode = BaseVar(
+    _var_name=constants.ColorMode.SET,
+    _var_type=EventChain,
+    _var_data=VarData(imports=color_mode_imports, hooks=color_mode_set_hooks),
+)
 
 breakpoints = ["0", "30em", "48em", "62em", "80em", "96em"]
 
@@ -273,7 +286,7 @@ def format_as_emotion(style_dict: dict[str, Any]) -> Style | None:
 
 
 def convert_dict_to_style_and_format_emotion(
-    raw_dict: dict[str, Any]
+    raw_dict: dict[str, Any],
 ) -> dict[str, Any] | None:
     """Convert a dict to a style dict and then format as emotion.