77import dataclasses
88import functools
99from pathlib import Path
10- from typing import Any , Dict
10+ from typing import Any , TypedDict
1111
1212import uproot
1313import uproot .reading
1818
1919console = Console ()
2020
21- __all__ = ("make_tree" , "process_item" , "print_tree" , "UprootEntry" , "console" )
21+ __all__ = (
22+ "make_tree" ,
23+ "process_item" ,
24+ "print_tree" ,
25+ "UprootEntry" ,
26+ "console" ,
27+ "MetaDict" ,
28+ )
2229
2330
2431def __dir__ () -> tuple [str , ...]:
2532 return __all__
2633
2734
35+ class MetaDictRequired (TypedDict , total = True ):
36+ label_text : Text
37+ label_icon : str
38+
39+
40+ class MetaDict (MetaDictRequired , total = False ):
41+ guide_style : str
42+
43+
2844@dataclasses .dataclass
2945class UprootEntry :
3046 path : str
@@ -38,11 +54,18 @@ def is_dir(self) -> bool:
3854 return len (self .item .branches ) > 0
3955 return False
4056
41- def meta (self ) -> dict [ str , Any ] :
57+ def meta (self ) -> MetaDict :
4258 return process_item (self .item )
4359
4460 def label (self ) -> Text :
45- return process_item (self .item )["label" ] # type: ignore[no-any-return]
61+ meta = self .meta ()
62+ return Text .assemble (meta ["label_icon" ], meta ["label_text" ])
63+
64+ def _tree_args (self ) -> dict [str , Any ]:
65+ d : dict [str , Text | str ] = {"label" : self .label ()}
66+ if "guide_style" in self .meta ():
67+ d ["guide_style" ] = self .meta ()["guide_style" ]
68+ return d
4669
4770 @property
4871 def children (self ) -> list [UprootEntry ]:
@@ -68,7 +91,7 @@ def make_tree(node: UprootEntry, *, tree: Tree | None = None) -> Tree:
6891 Given an object, build a rich.tree.Tree output.
6992 """
7093
71- tree = Tree (** node .meta ()) if tree is None else tree .add (** node .meta ())
94+ tree = Tree (** node ._tree_args ()) if tree is None else tree .add (** node ._tree_args ())
7295
7396 for child in node .children :
7497 make_tree (child , tree = tree )
@@ -77,24 +100,23 @@ def make_tree(node: UprootEntry, *, tree: Tree | None = None) -> Tree:
77100
78101
79102@functools .singledispatch
80- def process_item (uproot_object : Any ) -> Dict [ str , Any ] :
103+ def process_item (uproot_object : Any ) -> MetaDict :
81104 """
82105 Given an unknown object, return a rich.tree.Tree output. Specialize for known objects.
83106 """
84107 name = getattr (uproot_object , "name" , "<unnamed>" )
85108 classname = getattr (uproot_object , "classname" , uproot_object .__class__ .__name__ )
86- label = Text .assemble (
87- "❓ " ,
109+ label_text = Text .assemble (
88110 (f"{ name } " , "bold" ),
89111 (classname , "italic" ),
90112 )
91- return { "label" : label }
113+ return MetaDict ( label_icon = "❓ " , label_text = label_text )
92114
93115
94116@process_item .register
95117def _process_item_tfile (
96118 uproot_object : uproot .reading .ReadOnlyDirectory ,
97- ) -> Dict [ str , Any ] :
119+ ) -> MetaDict :
98120 """
99121 Given an TFile, return a rich.tree.Tree output.
100122 """
@@ -109,33 +131,34 @@ def _process_item_tfile(
109131 path_name = escape (path .name )
110132 link_text = f"file://{ path } "
111133
112- label = Text .from_markup (f":file_folder: [link { link_text } ]{ path_name } " )
134+ label_text = Text .from_markup (f"[link { link_text } ]{ path_name } " )
113135
114- return {
115- "label" : label ,
116- "guide_style" : "bold bright_blue" ,
117- }
136+ return MetaDict (
137+ label_icon = "📁 " ,
138+ label_text = label_text ,
139+ guide_style = "bold bright_blue" ,
140+ )
118141
119142
120143@process_item .register
121- def _process_item_ttree (uproot_object : uproot .TTree ) -> Dict [ str , Any ] :
144+ def _process_item_ttree (uproot_object : uproot .TTree ) -> MetaDict :
122145 """
123146 Given an tree, return a rich.tree.Tree output.
124147 """
125- label = Text .assemble (
126- "🌴 " ,
148+ label_text = Text .assemble (
127149 (f"{ uproot_object .name } " , "bold" ),
128150 f"({ uproot_object .num_entries :g} )" ,
129151 )
130152
131- return {
132- "label" : label ,
133- "guide_style" : "bold bright_green" ,
134- }
153+ return MetaDict (
154+ label_icon = "🌴 " ,
155+ label_text = label_text ,
156+ guide_style = "bold bright_green" ,
157+ )
135158
136159
137160@process_item .register
138- def _process_item_tbranch (uproot_object : uproot .TBranch ) -> Dict [ str , Any ] :
161+ def _process_item_tbranch (uproot_object : uproot .TBranch ) -> MetaDict :
139162 """
140163 Given an branch, return a rich.tree.Tree output.
141164 """
@@ -148,29 +171,35 @@ def _process_item_tbranch(uproot_object: uproot.TBranch) -> Dict[str, Any]:
148171 if len (uproot_object .branches ):
149172 icon = "🌿 "
150173
151- label = Text .assemble (
152- icon ,
174+ label_text = Text .assemble (
153175 (f"{ uproot_object .name } " , "bold" ),
154176 (f"{ uproot_object .typename } " , "italic" ),
155177 )
156- return {"label" : label }
178+
179+ return MetaDict (
180+ label_icon = icon ,
181+ label_text = label_text ,
182+ guide_style = "bold bright_green" ,
183+ )
157184
158185
159186@process_item .register
160- def _process_item_th (uproot_object : uproot .behaviors .TH1 .Histogram ) -> Dict [ str , Any ] :
187+ def _process_item_th (uproot_object : uproot .behaviors .TH1 .Histogram ) -> MetaDict :
161188 """
162189 Given an histogram, return a rich.tree.Tree output.
163190 """
164191 icon = "📊 " if uproot_object .kind == "COUNT" else "📈 "
165192 sizes = " × " .join (f"{ len (ax )} " for ax in uproot_object .axes )
166193
167- label = Text .assemble (
168- icon ,
194+ label_text = Text .assemble (
169195 (f"{ uproot_object .name } " , "bold" ),
170196 (f"{ uproot_object .classname } " , "italic" ),
171197 f"({ sizes } )" ,
172198 )
173- return {"label" : label }
199+ return MetaDict (
200+ label_icon = icon ,
201+ label_text = label_text ,
202+ )
174203
175204
176205# pylint: disable-next=redefined-outer-name
0 commit comments