diff --git a/src/uproot_browser/plot.py b/src/uproot_browser/plot.py index ed9cec9..cbc6d1d 100644 --- a/src/uproot_browser/plot.py +++ b/src/uproot_browser/plot.py @@ -13,6 +13,8 @@ import numpy as np import plotext as plt import uproot +import uproot.behaviors.TH1 +import uproot.models.RNTuple from uproot_browser.exceptions import EmptyTreeError @@ -42,7 +44,7 @@ def make_hist_title(item: Any, histogram: hist.Hist) -> str: @functools.singledispatch -def plot(tree: Any, *, expr: str = "") -> None: # noqa: ARG001 +def plot(tree: Any, *, width: int = 100, expr: str = "") -> None: # noqa: ARG001 """ Implement this for each type of plottable. """ @@ -53,7 +55,10 @@ def plot(tree: Any, *, expr: str = "") -> None: # noqa: ARG001 # Simpler in Python 3.11+ @plot.register(uproot.TBranch) def plot_branch( - tree: uproot.TBranch | uproot.models.RNTuple.RField, *, expr: str = "" + tree: uproot.TBranch | uproot.models.RNTuple.RField, + *, + width: int = 100, + expr: str = "", ) -> None: """ Plot a single tree branch. @@ -64,14 +69,13 @@ def plot_branch( if len(finite) < 1: msg = f"Branch {tree.name} is empty." raise EmptyTreeError(msg) - histogram: hist.Hist = hist.numpy.histogram(finite, bins=100, histogram=hist.Hist) + histogram: hist.Hist = hist.numpy.histogram(finite, bins=width, histogram=hist.Hist) if expr: # pylint: disable-next=eval-used histogram = eval(expr, {"h": histogram}) plt.bar( - histogram.axes[0].centers, + histogram.axes[0].edges, histogram.values().astype(float), - width=histogram.axes[0].widths, ) plt.ylim(lower=0) plt.xticks(np.linspace(histogram.axes[0][0][0], histogram.axes[0][-1][-1], 5)) @@ -83,7 +87,11 @@ def plot_branch( @plot.register -def plot_hist(tree: uproot.behaviors.TH1.Histogram, expr: str = "") -> None: +def plot_hist( + tree: uproot.behaviors.TH1.Histogram, + width: int = 100, # noqa: ARG001 + expr: str = "", +) -> None: """ Plot a 1-D Histogram. """ @@ -91,7 +99,7 @@ def plot_hist(tree: uproot.behaviors.TH1.Histogram, expr: str = "") -> None: if expr: # pylint: disable-next=eval-used histogram = eval(expr, {"h": histogram}) - plt.bar(histogram.axes[0].centers, histogram.values().astype(float)) + plt.bar(histogram.axes[0].edges, histogram.values().astype(float)) plt.ylim(lower=0) plt.xticks(np.linspace(histogram.axes[0][0][0], histogram.axes[0][-1][-1], 5)) plt.xlabel(histogram.axes[0].name) diff --git a/src/uproot_browser/tui/browser.py b/src/uproot_browser/tui/browser.py index 48aabe2..9d20f56 100644 --- a/src/uproot_browser/tui/browser.py +++ b/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 @@ -140,6 +141,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 {"", "__main__"}: fname = "../scikit-hep-testdata/src/skhep_testdata/data/uproot-Event.root" diff --git a/src/uproot_browser/tui/messages.py b/src/uproot_browser/tui/messages.py index 904e4e3..62534c8 100644 --- a/src/uproot_browser/tui/messages.py +++ b/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__() diff --git a/src/uproot_browser/tui/plot.py b/src/uproot_browser/tui/plot.py index 6a743f4..46a8b8a 100644 --- a/src/uproot_browser/tui/plot.py +++ b/src/uproot_browser/tui/plot.py @@ -7,13 +7,12 @@ import plotext as plt # plots in text import rich.text -import textual.widgets import uproot_browser.plot 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 @@ -32,7 +31,7 @@ def make_plot(item: Any, theme: str, *size: int, expr: str) -> Any: plt.clf() plt.theme(theme) plt.plotsize(*size) - uproot_browser.plot.plot(item, expr=expr) + uproot_browser.plot.plot(item, width=(size[0] - 5) * 4, expr=expr) return plt.build() @@ -44,6 +43,24 @@ class Plotext: theme: str app: Browser expr: str = "" + size: tuple[int, int] | None = None + previous: rich.text.Text | None = None + old_expr: str = "" + + 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, expr=self.expr) + 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 @@ -59,13 +76,17 @@ 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, expr=self.expr) - yield rich.text.Text.from_ansi(canvas) - except EmptyTreeError: - self.app.post_message(EmptyMessage()) - except Exception: - self.app.query_one("#plot-input", textual.widgets.Input).value = "" - 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 + and self.old_expr == self.expr + ): + yield self.previous + + else: + self.size = (width, height) + self.previous = rich.text.Text("... plotting ...") + self.old_expr = self.expr + yield self.previous + self.app.post_message(RequestPlot(self)) diff --git a/tests/test_tui.py b/tests/test_tui.py index 933b81c..a6720b8 100644 --- a/tests/test_tui.py +++ b/tests/test_tui.py @@ -16,6 +16,7 @@ async def test_browse_plot() -> None: skhep_testdata.data_path("uproot-Event.root") ).run_test() as pilot: await pilot.press("down", "down", "down", "enter") + await pilot.pause() assert isinstance(pilot.app.view_widget.item, Plotext) @@ -24,6 +25,7 @@ async def test_browse_empty() -> None: skhep_testdata.data_path("uproot-empty.root") ).run_test() as pilot: await pilot.press("down", "space", "down", "enter") + await pilot.pause() assert pilot.app.view_widget.item is None @@ -32,6 +34,7 @@ async def test_browse_empty_vim() -> None: skhep_testdata.data_path("uproot-empty.root") ).run_test() as pilot: await pilot.press("j", "l", "j", "enter") + await pilot.pause() assert pilot.app.view_widget.item is None