fix: better RNTuple support (#211)

* fix: better RNTuple support

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

* fix: bump minimum versions

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

---------

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

Author: henryiii

noxfile.py

@@ -6,10 +6,6 @@
 nox.options.default_venv_backend = "uv|virtualenv"
 
 
-def dep_group(group: str) -> list[str]:
-    return nox.project.load_toml("pyproject.toml")["dependency-groups"][group]  # type: ignore[no-any-return]
-
-
 @nox.session(reuse_venv=True)
 def lint(session: nox.Session) -> None:
     """
@@ -34,7 +30,7 @@ def tests(session: nox.Session) -> None:
     """
     Run the unit and regular tests.
     """
-    session.install("-e.", *dep_group("test"))
+    session.install("-e.", "--group=test")
     session.run("pytest", *session.posargs)
 
 
@@ -44,7 +40,7 @@ def minimums(session: nox.Session) -> None:
     Run the unit and regular tests.
     """
     session.install(
-        "-e.", *dep_group("test"), "--resolution=lowest-direct", "--only-binary=:all:"
+        "-e.", "--group=test", "--resolution=lowest-direct", "--only-binary=:all:"
     )
     session.run("pytest", *session.posargs)
 

pyproject.toml

@@ -45,7 +45,7 @@ dependencies = [
   'plotext >=5.2.8',
   'rich >=13.3.3',
   'textual >=0.86.0',
-  'uproot >=5',
+  'uproot >=5.6.1',
 ]
 
 [project.optional-dependencies]
@@ -55,7 +55,7 @@ iterm = [
   "mplhep",
 ]
 testdata = [
-  "scikit-hep-testdata >=0.4.10",
+  "scikit-hep-testdata >=0.5.0",
 ]
 
 [project.urls]
@@ -69,7 +69,7 @@ uproot-browser = "uproot_browser.__main__:main"
 test = [
   "pytest >=8",
   "pytest-asyncio >=0.24",
-  "scikit-hep-testdata >=0.4.10",
+  "scikit-hep-testdata >=0.5.0",
 ]
 dev = [
   "ipython >=6",

src/uproot_browser/plot.py

@@ -70,8 +70,7 @@ def plot_branch(tree: uproot.TBranch | uproot.models.RNTuple.RField) -> 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(uproot.models.RNTuple.RField)(plot_branch)  # type: ignore[no-untyped-call]
 
 
 @plot.register

src/uproot_browser/tree.py

@@ -57,13 +57,9 @@ 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
+@is_dir.register
+def _(item: uproot.behaviors.RNTuple.HasFields) -> bool:
+    return len(item.keys()) > 0
 
 
 def get_children(item: Mapping[str, Any]) -> set[str]:
@@ -161,11 +157,8 @@ def _process_item_tfile(
     )
 
 
-# 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:
+@process_item.register
+def _process_item_ttree(uproot_object: uproot.TTree) -> MetaDict:
     """
     Given an tree, return a rich.tree.Tree output.
     """
@@ -181,10 +174,23 @@ def _process_item_ttree(
     )
 
 
-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_rntuple(
+    uproot_object: uproot.behaviors.RNTuple.RNTuple,
+) -> MetaDict:
+    """
+    Given an tree, return a rich.tree.Tree output.
+    """
+    label_text = Text.assemble(
+        (f"{uproot_object.name} ", "bold"),
+        f"({uproot_object.num_entries:g})",
+    )
+
+    return MetaDict(
+        label_icon="🌳 ",
+        label_text=label_text,
+        guide_style="bold bright_green",
+    )
 
 
 @process_item.register
@@ -213,6 +219,26 @@ def _process_item_tbranch(uproot_object: uproot.TBranch) -> MetaDict:
     )
 
 
+@process_item.register
+def _process_item_rbranch(uproot_object: uproot.models.RNTuple.RField) -> MetaDict:
+    """
+    Given an branch, return a rich.tree.Tree output.
+    """
+
+    icon = "🍁 "
+
+    label_text = Text.assemble(
+        (f"{uproot_object.name} ", "bold"),
+        (f"{uproot_object.typename}", "italic"),
+    )
+
+    return MetaDict(
+        label_icon=icon,
+        label_text=label_text,
+        guide_style="bold bright_green",
+    )
+
+
 @process_item.register
 def _process_item_th(uproot_object: uproot.behaviors.TH1.Histogram) -> MetaDict:
     """

src/uproot_browser/tui/plot.py

@@ -46,7 +46,9 @@ class Plotext:
     def __rich_console__(
         self, console: rich.console.Console, options: rich.console.ConsoleOptions
     ) -> rich.console.RenderResult:
-        *_, item = apply_selection(self.upfile, self.selection.split(":"))
+        *_, item = apply_selection(
+            self.upfile, [s for s in self.selection.split("/") if s]
+        )
 
         if item is None:
             self.app.post_message(EmptyMessage())

tests/test_printouts.py

@@ -90,3 +90,34 @@ def test_tree(capsys):
 
     assert not err
     assert out == OUT1
+
+
+OUT2 = """\
+📁 ntpl001_staff_rntuple_v1-0-0-0.root
+┗━━ 🌳 Staff (3354)
+    ┣━━ 🍁 Age std::int32_t
+    ┣━━ 🍁 Category std::int32_t
+    ┣━━ 🍁 Children std::int32_t
+    ┣━━ 🍁 Cost std::int32_t
+    ┣━━ 🍁 Division std::string
+    ┣━━ 🍁 Flag std::uint32_t
+    ┣━━ 🍁 Grade std::int32_t
+    ┣━━ 🍁 Hrweek std::int32_t
+    ┣━━ 🍁 Nation std::string
+    ┣━━ 🍁 Service std::int32_t
+    ┗━━ 🍁 Step std::int32_t
+"""
+
+
+@pytest.mark.xfail(
+    sys.platform.startswith("win"),
+    reason="Unicode is different on Windows, for some reason?",
+)
+def test_tree_rntuple(capsys):
+    filename = data_path("ntpl001_staff_rntuple_v1-0-0-0.root")
+    console = rich.console.Console(width=120)
+
+    print_tree(filename, console=console)
+    out, err = capsys.readouterr()
+    assert not err
+    assert out == OUT2