diff --git a/src/uproot_browser/plot.py b/src/uproot_browser/plot.py index bb5fc1b..52a3724 100644 --- a/src/uproot_browser/plot.py +++ b/src/uproot_browser/plot.py @@ -50,8 +50,9 @@ def plot(tree: Any) -> None: raise RuntimeError(msg) -@plot.register -def plot_branch(tree: uproot.TBranch) -> None: +# Simpler in Python 3.11+ +@plot.register(uproot.TBranch) +def plot_branch(tree: uproot.TBranch | uproot.models.RNTuple.RField) -> None: """ Plot a single tree branch. """ @@ -69,6 +70,10 @@ def plot_branch(tree: uproot.TBranch) -> None: plt.title(make_hist_title(tree, histogram)) +if hasattr(uproot.models, "RNTuple") and hasattr(uproot.models.RNTuple, "RField"): + plot.register(uproot.models.RNTuple.RField)(plot_branch) # type: ignore[no-untyped-call] + + @plot.register def plot_hist(tree: uproot.behaviors.TH1.Histogram) -> None: """ diff --git a/src/uproot_browser/tree.py b/src/uproot_browser/tree.py index 946a9c6..414d77b 100644 --- a/src/uproot_browser/tree.py +++ b/src/uproot_browser/tree.py @@ -57,6 +57,15 @@ def _(item: uproot.behaviors.TBranch.HasBranches) -> bool: return len(item.branches) > 0 +if hasattr(uproot.behaviors, "RNTuple") and hasattr( + uproot.behaviors.RNTuple, "HasFields" +): + + @is_dir.register + def _(item: uproot.behaviors.RNTuple.HasFields) -> bool: + return len(item.keys()) > 0 + + def get_children(item: Mapping[str, Any]) -> set[str]: return { key.split(";")[0] @@ -152,8 +161,11 @@ def _process_item_tfile( ) -@process_item.register -def _process_item_ttree(uproot_object: uproot.TTree) -> MetaDict: +# Python 3.11 can just use `|` directly for register +@process_item.register(uproot.TTree) +def _process_item_ttree( + uproot_object: uproot.TTree | uproot.behaviors.RNTuple.RNTuple, +) -> MetaDict: """ Given an tree, return a rich.tree.Tree output. """ @@ -169,6 +181,12 @@ def _process_item_ttree(uproot_object: uproot.TTree) -> MetaDict: ) +if hasattr(uproot.behaviors, "RNTuple") and hasattr( + uproot.behaviors.RNTuple, "HasFields" +): + process_item.register(uproot.behaviors.RNTuple.RNTuple)(_process_item_ttree) # type: ignore[no-untyped-call] + + @process_item.register def _process_item_tbranch(uproot_object: uproot.TBranch) -> MetaDict: """