Skip to content

Commit 753159a

Browse files
ariostashenryiii
andauthored
fix: RNTuple tree structure (#216)
* Fixed RNTuple tree structure and minor refactoring * Update src/uproot_browser/tree.py --------- Co-authored-by: Henry Schreiner <[email protected]>
1 parent 5013a8c commit 753159a

File tree

1 file changed

+14
-16
lines changed

1 file changed

+14
-16
lines changed

src/uproot_browser/tree.py

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66

77
import dataclasses
88
import functools
9-
from collections.abc import Mapping
9+
from collections.abc import Iterable
1010
from pathlib import Path
11-
from typing import Any, Literal, TypedDict
11+
from typing import Any, Literal, Protocol, TypedDict
1212

1313
import uproot
1414
import uproot.reading
@@ -29,6 +29,10 @@
2929
)
3030

3131

32+
class SupportsRecursiveKeys(Protocol):
33+
def keys(self, recursive: bool = True) -> Iterable[str]: ...
34+
35+
3236
def __dir__() -> tuple[str, ...]:
3337
return __all__
3438

@@ -52,22 +56,16 @@ def _(item: uproot.reading.ReadOnlyDirectory) -> Literal[True]: # noqa: ARG001
5256
return True
5357

5458

55-
@is_dir.register
56-
def _(item: uproot.behaviors.TBranch.HasBranches) -> bool:
57-
return len(item.branches) > 0
58-
59-
60-
@is_dir.register
61-
def _(item: uproot.behaviors.RNTuple.HasFields) -> bool:
62-
return len(item.keys()) > 0
59+
@is_dir.register(uproot.behaviors.TBranch.HasBranches)
60+
@is_dir.register(uproot.behaviors.RNTuple.HasFields)
61+
def _(
62+
item: uproot.behaviors.TBranch.HasBranches | uproot.behaviors.RNTuple.HasFields,
63+
) -> bool:
64+
return len(item) > 0
6365

6466

65-
def get_children(item: Mapping[str, Any]) -> set[str]:
66-
return {
67-
key.split(";")[0]
68-
for key in item.keys() # noqa: SIM118
69-
if "/" not in key
70-
}
67+
def get_children(item: SupportsRecursiveKeys) -> set[str]:
68+
return {key.split(";")[0] for key in item.keys(recursive=False)}
7169

7270

7371
@dataclasses.dataclass

0 commit comments

Comments
 (0)