Add more docs, ensure contextmanager isn't reentrant

plugins/ui/src/deephaven/ui/_internal/RenderContext.py

@@ -49,22 +49,32 @@ class NoContextException(Exception):
 
 
 def get_context() -> RenderContext:
+    """
+    Gets the currently active context, or throws NoContextException if none is set.
+
+    Returns:
+        The active RenderContext, or throws if none is present.
+    """
     try:
         return _local_data.context
     except AttributeError:
         raise NoContextException("No context set")
 
 
-def set_context(context: Optional[RenderContext]):
+def _set_context(context: Optional[RenderContext]):
     """
-    Set the current context for the thread. Can be set to None to unset the context for a thread
+    Set the current context for the thread. Can be set to None to unset the context for a thread.
     """
     if context is None:
         del _local_data.context
     else:
         _local_data.context = context
 
 
+_READY_TO_OPEN: int = -2
+_OPENED_AND_UNUSED: int = -1
+
+
 class RenderContext:
     """
     Context for rendering a component. Keeps track of state and child contexts.
@@ -109,7 +119,7 @@ def __init__(self, on_change: OnChangeCallable, on_queue_render: OnChangeCallabl
             on_queue_render: The callback to call when work is being requested for the render loop.
         """
 
-        self._hook_index = -1
+        self._hook_index = _READY_TO_OPEN
         self._hook_count = -1
         self._state = {}
         self._children_context = {}
@@ -121,32 +131,42 @@ def __init__(self, on_change: OnChangeCallable, on_queue_render: OnChangeCallabl
     def open(self) -> AbstractContextManager:
         """
         Opens this context to track hook creation, sets this context as active on
-        this thread, and opens the liveness scope for user-created objects
-        :return:
+        this thread, and opens the liveness scope for user-created objects.
+
+        This is not reentrant and not safe across threads, ensure it is only opened
+        once at a time.
+
+        Returns:
+            A context manager to manage RenderContext resources.
         """
-        self._hook_index = -1
+        if self._hook_index != _READY_TO_OPEN:
+            raise RuntimeError(
+                "RenderContext.open() was already called, and is not reentrant"
+            )
+        self._hook_index = _OPENED_AND_UNUSED
 
         old_context: Optional[RenderContext] = None
         try:
             old_context = get_context()
         except NoContextException:
             pass
         logger.debug("old context is %s and new context is %s", old_context, self)
-        set_context(self)
+        _set_context(self)
 
         try:
             new_liveness_scope = LivenessScope()
             with new_liveness_scope.open():
                 yield self
 
             # Following the "yield" so we don't do this if there was an error, we want to keep the old scope and kill
-            # the new one
+            # the new one. We always release after creating the new one, so that each table/etc has its ref count go
+            # from 1 -> 2 -> 1, instead of 1 -> 0 -> 1 which would release the table prematurely
             self._liveness_scope.release()
             self._liveness_scope = new_liveness_scope
         finally:
             # Do this even if there was an error, old context must be restored
             logger.debug("Resetting to old context %s", old_context)
-            set_context(old_context)
+            _set_context(old_context)
 
         # Outside the "finally" so we don't do this if there was an error, we don't want an incorrect hook count
         hook_count = self._hook_index + 1
@@ -159,6 +179,9 @@ def open(self) -> AbstractContextManager:
                 )
             )
 
+        # Reset count for next use to safeguard double-opening
+        self._hook_index = _READY_TO_OPEN
+
     def has_state(self, key: StateKey) -> bool:
         """
         Check if the given key is in the state.

plugins/ui/test/deephaven/ui/test_render.py

@@ -14,7 +14,7 @@ class RenderTestCase(BaseTestCase):
     def test_empty_render(self):
         on_change = Mock(side_effect=lambda x: x())
         rc = make_render_context(on_change)
-        self.assertEqual(rc._hook_index, -1)
+        self.assertEqual(rc._hook_index, -2)
         self.assertEqual(rc._state, {})
         self.assertEqual(rc._children_context, {})
         on_change.assert_not_called()