fix: Stock rollup example was not working (#268)

- Referenced a table `t` that did not exist, corrected to reference
`formatted_table`
- Needed to refactor how exported objects were being handled - server
was just using a WeakKeyDictionary, but the objects were not guaranteed
to get cleared out on the server and were being reused between
re-renders. The client would always dump old objects after a render, so
just always clear out the objects on the server from the previous
render. Note that the WeakKeyDictionary is still safe for callables; the
client just references them by ID and does not hold a reference to an
object, so they can be cleared by the server GC.
- Fixes #265
- Tested using the Stock Rollup example fixed in this PR, it now works
correctly.
- Also tested with snippet from #128, ensuring the table still does not
flicker.
- Added another test to test more than one iteration

plugins/ui/examples/README.md

@@ -807,7 +807,12 @@ def stock_table(source):
         [source, highlight],
     )
     rolled_table = use_memo(
-        lambda: t if len(by) == 0 else t.rollup(aggs=aggs, by=by), [t, aggs, by]
+        lambda: (
+            formatted_table
+            if len(by) == 0
+            else formatted_table.rollup(aggs=aggs, by=by)
+        ),
+        [formatted_table, aggs, by],
     )
 
     return ui.flex(
@@ -816,20 +821,22 @@ def stock_table(source):
             ui.toggle_button(
                 ui.icon("vsBell"), "By Exchange", on_change=set_is_exchange
             ),
-            ui.fragment(
-                ui.text_field(
-                    label="Highlight Sym",
-                    label_position="side",
-                    value=highlight,
-                    on_change=set_highlight,
-                ),
-                ui.contextual_help(
-                    ui.heading("Highlight Sym"),
-                    ui.content("Enter a sym you would like highlighted."),
-                ),
-            )
-            if not is_sym and not is_exchange
-            else None,
+            (
+                ui.fragment(
+                    ui.text_field(
+                        label="Highlight Sym",
+                        label_position="side",
+                        value=highlight,
+                        on_change=set_highlight,
+                    ),
+                    ui.contextual_help(
+                        ui.heading("Highlight Sym"),
+                        ui.content("Enter a sym you would like highlighted."),
+                    ),
+                )
+                if not is_sym and not is_exchange
+                else None
+            ),
             align_items="center",
             gap="size-100",
             margin="size-100",

plugins/ui/src/deephaven/ui/renderer/NodeEncoder.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any, Callable, TypedDict
 from weakref import WeakKeyDictionary
 from .RenderedNode import RenderedNode
 
+logger = logging.getLogger(__name__)
+
 CALLABLE_KEY = "__dhCbid"
 OBJECT_KEY = "__dhObid"
 ELEMENT_KEY = "__dhElemName"
@@ -20,7 +23,7 @@
 
 class NodeEncoderResult(TypedDict):
     """
-    Result of encoding a node. Contains the encoded node, list of new objects, and callables dictionary.
+    Result of encoding a node. Contains the encoded node, set of new objects, and callables dictionary.
     """
 
     encoded_node: str
@@ -72,9 +75,15 @@ class NodeEncoder(json.JSONEncoder):
     The next object id to use. Increment for each new object encountered.
     """
 
-    _object_id_dict: WeakKeyDictionary[Any, int]
+    _old_objects: set[Any]
+    """
+    List of objects from the last render. Used to remove objects that are no longer in the document.
+    """
+
+    _object_id_dict: dict[Any, int]
     """
-    Dictionary from an object to the ID assigned to it
+    Dictionary from an object to the ID assigned to it. Objects are removed after the next render if they are no longer in the document.
+    Unlike `_callable_dict`, we cannot use a WeakKeyDictionary as we need to pass the exported object instance to the client, so we need to always keep a reference around that the client may still have a reference to.
     """
 
     def __init__(
@@ -97,7 +106,7 @@ def __init__(
         self._callable_dict = WeakKeyDictionary()
         self._new_objects = []
         self._next_object_id = 0
-        self._object_id_dict = WeakKeyDictionary()
+        self._object_id_dict = {}
 
     def default(self, o: Any):
         if isinstance(o, RenderedNode):
@@ -122,9 +131,22 @@ def encode_node(self, node: RenderedNode) -> NodeEncoderResult:
         Args:
             o: The document to encode
         """
+
         # Reset the new objects list - they will get set when encoding
         self._new_objects = []
+        self._old_objects = set(self._object_id_dict.keys())
+
+        logger.debug("Encoding node with object_id_dict: %s", self._object_id_dict)
+
         encoded_node = super().encode(node)
+
+        # Remove the old objects from last render from the object id dict
+        for obj in self._old_objects:
+            self._object_id_dict.pop(obj, None)
+
+        # Clear out the old objects list so they can be cleaned up by GC
+        self._old_objects = set()
+
         return {
             "encoded_node": encoded_node,
             "new_objects": self._new_objects,
@@ -156,6 +178,9 @@ def _convert_object(self, obj: Any):
             self._object_id_dict[obj] = object_id
             self._new_objects.append(obj)
 
+        self._old_objects.discard(obj)
+        logger.debug("Converted object %s to id %s", obj, object_id)
+
         return {
             OBJECT_KEY: object_id,
         }

plugins/ui/test/deephaven/ui/test_encoder.py

@@ -449,6 +449,7 @@ def test_delta_objects(self):
         )
 
         result = encoder.encode_node(node)
+
         expected_payload = {
             "__dhElemName": "test0",
             "props": {
@@ -491,6 +492,71 @@ def test_delta_objects(self):
         )
         self.assertListEqual(result["new_objects"], delta_objs, "objects don't match")
 
+        # Another step - Add some new objects and callables to the mix for next encoding
+        delta_objs = [TestObject()]
+        delta_cbs = [lambda: None]
+        objs = [objs[0], None, objs[2], None, delta_objs[0]]
+        cbs = [cbs[0], None, cbs[2], None, delta_cbs[0]]
+
+        # Replace cb[3] and obj[3] with the new callable/object and encode again
+        node = make_node(
+            "test0",
+            {
+                "children": [
+                    make_node("test1", {"foo": [cbs[0]]}),
+                    cbs[4],
+                    make_node("test3", {"bar": cbs[2], "children": [objs[0], objs[4]]}),
+                    make_node("test4", {"children": [objs[0], objs[2]]}),
+                    objs[4],
+                    make_node("test6", {"children": [objs[2]]}),
+                ]
+            },
+        )
+
+        result = encoder.encode_node(node)
+
+        expected_payload = {
+            "__dhElemName": "test0",
+            "props": {
+                "children": [
+                    {
+                        "__dhElemName": "test1",
+                        "props": {"foo": [{"__dhCbid": "cb0"}]},
+                    },
+                    {"__dhCbid": "cb4"},
+                    {
+                        "__dhElemName": "test3",
+                        "props": {
+                            "bar": {"__dhCbid": "cb2"},
+                            "children": [
+                                {"__dhObid": 0},
+                                {"__dhObid": 4},
+                            ],
+                        },
+                    },
+                    {
+                        "__dhElemName": "test4",
+                        "props": {"children": [{"__dhObid": 0}, {"__dhObid": 2}]},
+                    },
+                    {"__dhObid": 4},
+                    {
+                        "__dhElemName": "test6",
+                        "props": {"children": [{"__dhObid": 2}]},
+                    },
+                ],
+            },
+        }
+        self.maxDiff = None
+        self.assertDictEqual(
+            json.loads(result["encoded_node"]), expected_payload, "payloads don't match"
+        )
+        self.assertListEqual(
+            list(result["callable_id_dict"].keys()),
+            [cbs[0], cbs[2], cbs[4]],
+            "callables don't match",
+        )
+        self.assertListEqual(result["new_objects"], delta_objs, "objects don't match")
+
 
 if __name__ == "__main__":
     unittest.main()