fix: added support for autocompletion on the low ui approach for pycharm

src/streamsync/core.py

@@ -17,7 +17,7 @@
 import json
 import math
 from streamsync.ss_types import Readable, InstancePath, StreamsyncEvent, StreamsyncEventResult, StreamsyncFileItem
-from streamsync.core_ui import ComponentTree, SessionComponentTree
+from streamsync.core_ui import ComponentTree, SessionComponentTree, use_component_tree
 
 
 class Config:
@@ -1216,14 +1216,15 @@ def _call_handler_callable(self, event_type, target_component, instance_path, pa
                 arg_values.append(session_info)
             elif arg == "ui":
                 from streamsync.ui import StreamsyncUIManager
-                ui_manager = StreamsyncUIManager(self.session.session_component_tree)
+                ui_manager = StreamsyncUIManager()
                 arg_values.append(ui_manager)
 
         result = None
-        if is_async_handler:
-            result, captured_stdout = self._async_handler_executor(callable_handler, arg_values)
-        else:
-            result, captured_stdout = self._sync_handler_executor(callable_handler, arg_values)
+        with use_component_tree(self.session.session_component_tree):
+            if is_async_handler:
+                result, captured_stdout = self._async_handler_executor(callable_handler, arg_values)
+            else:
+                result, captured_stdout = self._sync_handler_executor(callable_handler, arg_values)
 
         if captured_stdout:
             self.session_state.add_log_entry(

src/streamsync/core_ui.py

@@ -1,11 +1,13 @@
+import contextlib
 from contextvars import ContextVar
 from typing import Any, Dict, List, Optional, Union
 import uuid
 
 from pydantic import BaseModel, Field
 
-current_parent_container: ContextVar[Union["Component", None]] = \
-    ContextVar("current_parent_container")
+
+current_parent_container: ContextVar[Union["Component", None]] = ContextVar("current_parent_container")
+_current_component_tree: ContextVar[Union["ComponentTree", None]] = ContextVar("current_component_tree", default=None)
 # This variable is thread safe and context safe
 
 
@@ -161,3 +163,35 @@ def fetch_updates(self):
 
 class UIError(Exception):
     ...
+
+@contextlib.contextmanager
+def use_component_tree(component_tree: ComponentTree):
+    """
+    Declares the component tree that will be manipulated during a context.
+
+    The declared tree can be retrieved with the `current_component_tree` method.
+
+    >>> with use_component_tree(component_tree):
+    >>>     ui_manager = StreamsyncUIManager()
+    >>>     ui_manager.create_component("text", text="Hello, world!")
+
+    :param component_tree:
+    """
+    token = _current_component_tree.set(component_tree)
+    yield
+    _current_component_tree.reset(token)
+
+
+def current_component_tree() -> ComponentTree:
+    """
+    Retrieves the component tree of the current context or the base
+    one if no context has been declared.
+
+    :return:
+    """
+    tree = _current_component_tree.get()
+    if tree is None:
+        from streamsync.core import base_component_tree
+        return base_component_tree
+
+    return tree