fix: faster (threaded) plotting

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

src/uproot_browser/tui/browser.py

@@ -13,6 +13,7 @@
 import textual.containers
 import textual.events
 import textual.widgets
+import textual.worker
 from textual.reactive import var
 
 with contextlib.suppress(AttributeError):
@@ -36,7 +37,7 @@
 from .header import Header
 from .help import HelpScreen
 from .left_panel import UprootTree
-from .messages import ErrorMessage, UprootSelected
+from .messages import ErrorMessage, RequestPlot, UprootSelected
 from .plot import Plotext
 from .tools import Info, Tools
 from .viewer import ViewWidget
@@ -138,6 +139,16 @@ def on_empty_message(self) -> None:
     def on_error_message(self, message: ErrorMessage) -> None:
         self.view_widget.item = message.err
 
+    def on_request_plot(self, message: RequestPlot) -> None:
+        self.render_plot(message.plot)
+
+    @textual.work(exclusive=True, thread=True)
+    def render_plot(self, plot: Plotext) -> None:
+        worker = textual.worker.get_current_worker()
+        new_plot = plot.make_plot()
+        if new_plot and not worker.is_cancelled:
+            self.call_from_thread(self.view_widget.plot_widget.update, new_plot)
+
 
 if __name__ in {"<run_path>", "__main__"}:
     fname = "../scikit-hep-testdata/src/skhep_testdata/data/uproot-Event.root"

src/uproot_browser/tui/messages.py

@@ -7,6 +7,7 @@
 
 if TYPE_CHECKING:
     from .error import Error
+    from .plot import Plotext
 
 
 @rich.repr.auto
@@ -27,3 +28,10 @@ class ErrorMessage(textual.message.Message, bubble=True):
     def __init__(self, err: Error) -> None:
         self.err = err
         super().__init__()
+
+
+@rich.repr.auto
+class RequestPlot(textual.message.Message, bubble=True):
+    def __init__(self, plot: Plotext) -> None:
+        self.plot = plot
+        super().__init__()

src/uproot_browser/tui/plot.py

@@ -12,7 +12,7 @@
 from uproot_browser.exceptions import EmptyTreeError
 
 from .error import Error
-from .messages import EmptyMessage, ErrorMessage
+from .messages import EmptyMessage, ErrorMessage, RequestPlot
 
 if TYPE_CHECKING:
     from .browser import Browser
@@ -42,6 +42,23 @@ class Plotext:
     selection: str
     theme: str
     app: Browser
+    size: tuple[int, int] | None = None
+    previous: rich.text.Text | None = None
+
+    def make_plot(self) -> Plotext | None:
+        *_, item = apply_selection(self.upfile, self.selection.split(":"))
+        assert self.size
+        try:
+            canvas = make_plot(item, self.theme, *self.size)
+            return dataclasses.replace(self, previous=rich.text.Text.from_ansi(canvas))
+        except EmptyTreeError:
+            self.app.post_message(EmptyMessage())
+            return None
+        except Exception:
+            exc = sys.exc_info()
+            assert exc[1]
+            self.app.post_message(ErrorMessage(Error(exc)))
+            return None
 
     def __rich_console__(
         self, console: rich.console.Console, options: rich.console.ConsoleOptions
@@ -55,12 +72,11 @@ def __rich_console__(
         width = options.max_width or console.width
         height = options.height or console.height
 
-        try:
-            canvas = make_plot(item, self.theme, width, height)
-            yield rich.text.Text.from_ansi(canvas)
-        except EmptyTreeError:
-            self.app.post_message(EmptyMessage())
-        except Exception:
-            exc = sys.exc_info()
-            assert exc[1]
-            self.app.post_message(ErrorMessage(Error(exc)))
+        if self.size and (width, height) == self.size and self.previous is not None:
+            yield self.previous
+
+        else:
+            self.size = (width, height)
+            self.previous = rich.text.Text("... plotting ...")
+            yield self.previous
+            self.app.post_message(RequestPlot(self))