From fdc21d507ebf6b362b8e6c50515b497101824a21 Mon Sep 17 00:00:00 2001
From: Iisakki Rotko <iisakki.rotko@widgetti.io>
Date: Mon, 11 Mar 2024 10:05:50 +0100
Subject: [PATCH] fix: solara theming not respecting theme.js file (#536)

Fixes https://github.com/widgetti/solara/issues/529.
---
 solara/lab/components/theming.py     | 24 +++++++++++++++++++++++-
 solara/server/app.py                 | 27 +++++++++++++++++++--------
 solara/server/static/main-vuetify.js |  2 +-
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/solara/lab/components/theming.py b/solara/lab/components/theming.py
index cb7e948c0..390810854 100644
--- a/solara/lab/components/theming.py
+++ b/solara/lab/components/theming.py
@@ -1,4 +1,4 @@
-from typing import Callable, cast
+from typing import Callable, Dict, Union, cast
 
 import ipyvuetify.Themes
 from ipyvuetify.Themes import Theme
@@ -11,6 +11,28 @@
 ipyvuetify.Themes.theme = cast(ipyvuetify.Themes.Theme, theme)
 
 
+def _set_theme(themes: Union[Dict[str, Dict[str, str]], None]):
+    if themes is None:
+        return
+
+    for theme_type in themes.keys():
+        widget = getattr(theme.themes, theme_type)
+        with widget.hold_trait_notifications():
+            for k, v in themes[theme_type].items():
+                setattr(widget, k, v)
+
+
+def _get_theme(theme: Theme) -> Dict[str, Dict[str, str]]:
+    theme_dict: Dict[str, Dict[str, str]] = cast(Dict[str, Dict[str, str]], {})
+    for theme_type, theme_value in theme.themes.__dict__.items():
+        theme_traits = theme_value.keys
+        theme_dict[theme_type] = {}
+        for trait in theme_traits:
+            if not trait.startswith("_"):
+                theme_dict[theme_type][trait] = getattr(theme_value, trait)
+    return theme_dict
+
+
 @component_vue("theming.vue")
 def _ThemeToggle(
     theme_dark: str,
diff --git a/solara/server/app.py b/solara/server/app.py
index 31ff174f8..43bf1bf97 100644
--- a/solara/server/app.py
+++ b/solara/server/app.py
@@ -361,6 +361,16 @@ def load_app_widget(app_state, app_script: AppScript, pathname: str):
             container.children = [widget]
 
 
+def load_themes(themes: Dict[str, Dict[str, str]], dark: bool):
+    # While these usually gets set from the frontend, in solara (server) we want to know theme information directly at the first
+    # render. Also, using the same trait allows us to write code which works on all widgets platforms, instead
+    # or using something different when running under solara server
+    from solara.lab.components.theming import _set_theme, theme
+
+    _set_theme(themes)
+    theme.dark_effective = dark
+
+
 def solara_comm_target(comm, msg_first):
     app: Optional[AppScript] = None
 
@@ -374,18 +384,13 @@ def on_msg(msg):
             app_name = args.get("appName") or "__default__"
             app = apps[app_name]
             context = kernel_context.get_current_context()
-            dark = args.get("dark", False)
             import ipyvuetify
 
-            from solara.lab import theme
-
-            # While this usually gets set from the frontend, in solara (server) we want to know this directly at the first
-            # render. Also, using the same trait allows us to write code which works on all widgets platforms, instead
-            # or using something different when running under solara server
-            theme.dark_effective = dark
-
             container = ipyvuetify.Html(tag="div")
             context.container = container
+            themes = args.get("themes")
+            dark = args.get("dark")
+            load_themes(themes, dark)
             load_app_widget(None, app, path)
             comm.send({"method": "finished", "widget_id": context.container._model_id})
         elif method == "app-status":
@@ -399,11 +404,17 @@ def on_msg(msg):
                 comm.send({"method": "app-status", "started": False})
 
         elif method == "reload":
+            from solara.lab.components.theming import _get_theme, theme
+
             assert app is not None
             context = kernel_context.get_current_context()
             path = data.get("path", "")
+            current_theme = theme._instance.value
+            theme_dict = _get_theme(current_theme)
+
             with context:
                 context.restart()
+                load_themes(theme_dict, current_theme.dark_effective)
                 load_app_widget(context.state, app, path)
                 comm.send({"method": "finished"})
         else:
diff --git a/solara/server/static/main-vuetify.js b/solara/server/static/main-vuetify.js
index d1f9c8c6d..640b32f0b 100644
--- a/solara/server/static/main-vuetify.js
+++ b/solara/server/static/main-vuetify.js
@@ -230,7 +230,7 @@ async function solaraInit(mountId, appName) {
     if (kernelId && widgetModelId) {
         await widgetManager.fetchAll();
     } else {
-        widgetModelId = await widgetManager.run(appName, {path, dark: inDarkMode()});
+        widgetModelId = await widgetManager.run(appName, {path, dark: inDarkMode(), themes: vuetifyThemes});
     }
     await solaraMount(widgetManager, mountId || 'content', widgetModelId);
     skipReconnectedCheck = false;