From 28f29e622fd9f6dc4cf5f3dde1d7bd14f0f95167 Mon Sep 17 00:00:00 2001
From: "Maarten A. Breddels" <maartenbreddels@gmail.com>
Date: Tue, 21 Nov 2023 21:32:26 +0100
Subject: [PATCH] perf: ipyvue widget can use a faster less generate state_get

---
 solara/server/patch.py | 58 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/solara/server/patch.py b/solara/server/patch.py
index 4e74baa8d..29b0b6b81 100644
--- a/solara/server/patch.py
+++ b/solara/server/patch.py
@@ -296,6 +296,58 @@ def Output_exit(self, exc_type, exc_value, traceback):
         ip.display_pub._hooks.pop()
 
 
+def patch_ipyvue_performance():
+    import functools
+    from collections.abc import Iterable
+
+    @functools.lru_cache(None)  # type: ignore
+    def class_traits(cls, **metadata):
+        # we cache it for performance reasons
+        return cls.class_traits(**metadata)
+
+    @functools.lru_cache(None)  # type: ignore
+    def to_jsons_meta(cls):
+        traits = class_traits(cls)
+        callables = {}
+        for name, trait in traits.items():
+            to_json = trait.metadata.get("to_json")
+            if to_json:
+                callables[name] = to_json
+        return callables
+
+    def get_state_fast(self, key=None, drop_defaults=False):
+        cls = type(self)
+        traits = class_traits(cls, sync=True)  # type: ignore
+        if key is None:
+            keys = list(traits)
+        elif isinstance(key, str):
+            keys = [key]
+        elif isinstance(key, Iterable):
+            keys = list(key)
+        else:
+            raise ValueError("key must be a string, an iterable of keys, or None")
+        state = {}
+        to_jsons = to_jsons_meta(cls)  # type: ignore
+        assert drop_defaults is False
+        trait_values = self._trait_values
+        for k in keys:
+            if k not in trait_values:
+                value = getattr(self, k)
+            else:
+                value = trait_values[k]
+            if k in to_jsons:
+                wire_value = to_jsons[k](value, self)
+            else:
+                # should we call _trait_to_json?
+                wire_value = value
+            state[k] = wire_value
+        return state
+
+    import ipyvue
+
+    ipyvue.VueWidget.get_state = get_state_fast
+
+
 def patch():
     global _patched
     global global_widgets_dict
@@ -305,6 +357,12 @@ def patch():
     _patched = True
     __builtins__["display"] = IPython.display.display
 
+    if settings.main.experimental_performance:
+        # this might be a bit too much
+        # import traitlets
+        # traitlets.TraitType._validate = lambda self, trait, value: value
+
+        patch_ipyvue_performance()
     # the ipyvue.Template module cannot be accessed like ipyvue.Template
     # because the import in ipvue overrides it
     template_mod = sys.modules["ipyvue.Template"]