Skip to content

Commit 2821b65

Browse files
refactor: more singledispatch (#197)
* refactor: more singledispatch Signed-off-by: Henry Schreiner <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: somethign Signed-off-by: Henry Schreiner <[email protected]> * fix: simplify children function Signed-off-by: Henry Schreiner <[email protected]> --------- Signed-off-by: Henry Schreiner <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent a2c6438 commit 2821b65

File tree

1 file changed

+29
-17
lines changed

1 file changed

+29
-17
lines changed

src/uproot_browser/tree.py

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66

77
import dataclasses
88
import functools
9+
from collections.abc import Mapping
910
from pathlib import Path
10-
from typing import Any, TypedDict
11+
from typing import Any, Literal, TypedDict
1112

1213
import uproot
1314
import uproot.reading
@@ -41,18 +42,37 @@ class MetaDict(MetaDictRequired, total=False):
4142
guide_style: str
4243

4344

45+
@functools.singledispatch
46+
def is_dir(item: Any) -> bool: # noqa: ARG001
47+
return False
48+
49+
50+
@is_dir.register
51+
def _(item: uproot.reading.ReadOnlyDirectory) -> Literal[True]: # noqa: ARG001
52+
return True
53+
54+
55+
@is_dir.register
56+
def _(item: uproot.behaviors.TBranch.HasBranches) -> bool:
57+
return len(item.branches) > 0
58+
59+
60+
def get_children(item: Mapping[str, Any]) -> set[str]:
61+
return {
62+
key.split(";")[0]
63+
for key in item.keys() # noqa: SIM118
64+
if "/" not in key
65+
}
66+
67+
4468
@dataclasses.dataclass
4569
class UprootEntry:
4670
path: str
4771
item: Any
4872

4973
@property
5074
def is_dir(self) -> bool:
51-
if isinstance(self.item, uproot.reading.ReadOnlyDirectory):
52-
return True
53-
if isinstance(self.item, uproot.behaviors.TBranch.HasBranches):
54-
return len(self.item.branches) > 0
55-
return False
75+
return is_dir(self.item)
5676

5777
def meta(self) -> MetaDict:
5878
return process_item(self.item)
@@ -71,18 +91,10 @@ def tree_args(self) -> dict[str, Any]:
7191
def children(self) -> list[UprootEntry]:
7292
if not self.is_dir:
7393
return []
74-
if isinstance(self.item, uproot.reading.ReadOnlyDirectory):
75-
items = {
76-
key.split(";")[0]
77-
for key in self.item.keys() # noqa: SIM118
78-
if "/" not in key
79-
}
80-
elif isinstance(self.item, uproot.behaviors.TBranch.HasBranches):
81-
items = {item.name for item in self.item.branches}
82-
else:
83-
items = {obj.name.split(";")[0] for obj in self.item.branches}
94+
8495
return [
85-
UprootEntry(f"{self.path}/{key}", self.item[key]) for key in sorted(items)
96+
UprootEntry(f"{self.path}/{key}", self.item[key])
97+
for key in sorted(get_children(self.item))
8698
]
8799

88100

0 commit comments

Comments
 (0)