Skip to content

Commit 2ed955a

Browse files
authored
feat: support RNTuple (#195)
* feat: support RNTuple Signed-off-by: Henry Schreiner <[email protected]> * fix: support Python before 3.11 Signed-off-by: Henry Schreiner <[email protected]> * fix: conditional RNtuple Signed-off-by: Henry Schreiner <[email protected]> --------- Signed-off-by: Henry Schreiner <[email protected]>
1 parent 2821b65 commit 2ed955a

File tree

2 files changed

+27
-4
lines changed

2 files changed

+27
-4
lines changed

src/uproot_browser/plot.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,9 @@ def plot(tree: Any) -> None:
5050
raise RuntimeError(msg)
5151

5252

53-
@plot.register
54-
def plot_branch(tree: uproot.TBranch) -> None:
53+
# Simpler in Python 3.11+
54+
@plot.register(uproot.TBranch)
55+
def plot_branch(tree: uproot.TBranch | uproot.models.RNTuple.RField) -> None:
5556
"""
5657
Plot a single tree branch.
5758
"""
@@ -69,6 +70,10 @@ def plot_branch(tree: uproot.TBranch) -> None:
6970
plt.title(make_hist_title(tree, histogram))
7071

7172

73+
if hasattr(uproot.models, "RNTuple") and hasattr(uproot.models.RNTuple, "RField"):
74+
plot.register(uproot.models.RNTuple.RField)(plot_branch) # type: ignore[no-untyped-call]
75+
76+
7277
@plot.register
7378
def plot_hist(tree: uproot.behaviors.TH1.Histogram) -> None:
7479
"""

src/uproot_browser/tree.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,15 @@ def _(item: uproot.behaviors.TBranch.HasBranches) -> bool:
5757
return len(item.branches) > 0
5858

5959

60+
if hasattr(uproot.behaviors, "RNTuple") and hasattr(
61+
uproot.behaviors.RNTuple, "HasFields"
62+
):
63+
64+
@is_dir.register
65+
def _(item: uproot.behaviors.RNTuple.HasFields) -> bool:
66+
return len(item.keys()) > 0
67+
68+
6069
def get_children(item: Mapping[str, Any]) -> set[str]:
6170
return {
6271
key.split(";")[0]
@@ -152,8 +161,11 @@ def _process_item_tfile(
152161
)
153162

154163

155-
@process_item.register
156-
def _process_item_ttree(uproot_object: uproot.TTree) -> MetaDict:
164+
# Python 3.11 can just use `|` directly for register
165+
@process_item.register(uproot.TTree)
166+
def _process_item_ttree(
167+
uproot_object: uproot.TTree | uproot.behaviors.RNTuple.RNTuple,
168+
) -> MetaDict:
157169
"""
158170
Given an tree, return a rich.tree.Tree output.
159171
"""
@@ -169,6 +181,12 @@ def _process_item_ttree(uproot_object: uproot.TTree) -> MetaDict:
169181
)
170182

171183

184+
if hasattr(uproot.behaviors, "RNTuple") and hasattr(
185+
uproot.behaviors.RNTuple, "HasFields"
186+
):
187+
process_item.register(uproot.behaviors.RNTuple.RNTuple)(_process_item_ttree) # type: ignore[no-untyped-call]
188+
189+
172190
@process_item.register
173191
def _process_item_tbranch(uproot_object: uproot.TBranch) -> MetaDict:
174192
"""

0 commit comments

Comments
 (0)