diff --git a/src/uproot_browser/tree.py b/src/uproot_browser/tree.py index 7025d5f..946a9c6 100644 --- a/src/uproot_browser/tree.py +++ b/src/uproot_browser/tree.py @@ -6,8 +6,9 @@ import dataclasses import functools +from collections.abc import Mapping from pathlib import Path -from typing import Any, TypedDict +from typing import Any, Literal, TypedDict import uproot import uproot.reading @@ -41,6 +42,29 @@ class MetaDict(MetaDictRequired, total=False): guide_style: str +@functools.singledispatch +def is_dir(item: Any) -> bool: # noqa: ARG001 + return False + + +@is_dir.register +def _(item: uproot.reading.ReadOnlyDirectory) -> Literal[True]: # noqa: ARG001 + return True + + +@is_dir.register +def _(item: uproot.behaviors.TBranch.HasBranches) -> bool: + return len(item.branches) > 0 + + +def get_children(item: Mapping[str, Any]) -> set[str]: + return { + key.split(";")[0] + for key in item.keys() # noqa: SIM118 + if "/" not in key + } + + @dataclasses.dataclass class UprootEntry: path: str @@ -48,11 +72,7 @@ class UprootEntry: @property def is_dir(self) -> bool: - if isinstance(self.item, uproot.reading.ReadOnlyDirectory): - return True - if isinstance(self.item, uproot.behaviors.TBranch.HasBranches): - return len(self.item.branches) > 0 - return False + return is_dir(self.item) def meta(self) -> MetaDict: return process_item(self.item) @@ -71,18 +91,10 @@ def tree_args(self) -> dict[str, Any]: def children(self) -> list[UprootEntry]: if not self.is_dir: return [] - if isinstance(self.item, uproot.reading.ReadOnlyDirectory): - items = { - key.split(";")[0] - for key in self.item.keys() # noqa: SIM118 - if "/" not in key - } - elif isinstance(self.item, uproot.behaviors.TBranch.HasBranches): - items = {item.name for item in self.item.branches} - else: - items = {obj.name.split(";")[0] for obj in self.item.branches} + return [ - UprootEntry(f"{self.path}/{key}", self.item[key]) for key in sorted(items) + UprootEntry(f"{self.path}/{key}", self.item[key]) + for key in sorted(get_children(self.item)) ]