Skip to content
This repository was archived by the owner on Jun 21, 2022. It is now read-only.

Commit 565e927

Browse files
jpivarskijpivarski
committed
nested directories
1 parent cd5a7b5 commit 565e927

File tree

5 files changed

+65
-20
lines changed

5 files changed

+65
-20
lines changed

‎tests/README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
| `std::vector`, `std::string` (???) | **(unsure)** |
1515
| branch with "speed bumps" | ✓ |
1616
| all compression algorithms (none, zlib, lzma, lz4; ignoring "old") | ✓ |
17-
| `TTree` versions from 16 (2009) to 19 (present) | ✓ |
18-
| nested directories | **no!** |
17+
| files from 2009 (`TTree` version 16) to present (`TTree` version 19) | ✓ |
18+
| nested directories, cycle numbers, '/' and ';' notation | ✓ |
1919
| arrays interface | ✓ |
2020
| iterator interface | ✓ |
2121
| selection by list of branch names | **no!** |

‎tests/nesteddirs.root

44.5 KB
Binary file not shown.

‎tests/test_tree.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,3 +218,28 @@ def test_tree_iterator4(self):
218218
for arrays in uproot.iterator(1, ["tests/foriter2.root", "tests/foriter2.root"], "foriter2"):
219219
self.assertEqual(arrays[b"data"].tolist(), words2[i:i + 1])
220220
i += 1
221+
222+
def test_directories(self):
223+
file = uproot.open("tests/nesteddirs.root")
224+
self.assertEqual(file.contents, {b"one;1": b"TDirectory", b"three;1": b"TDirectory"})
225+
self.assertEqual(file.allcontents, {b"one/two;1": b"TDirectory", b"one/two/tree;1": b"TTree", b"three/tree;1": b"TTree", b"one;1": b"TDirectory", b"one/tree;1": b"TTree", b"three;1": b"TDirectory"})
226+
227+
self.assertEqual(file["one"]["tree"].branchnames, [b"one", b"two", b"three"])
228+
self.assertEqual(file["one"].get("tree", 1).branchnames, [b"one", b"two", b"three"])
229+
self.assertEqual(file["one/tree;1"].branchnames, [b"one", b"two", b"three"])
230+
self.assertEqual(file["one/two/tree;1"].branchnames, [b"Int32", b"Int64", b"UInt32", b"UInt64", b"Float32", b"Float64", b"Str", b"ArrayInt32", b"ArrayInt64", b"ArrayUInt32", b"ArrayUInt64", b"ArrayFloat32", b"ArrayFloat64", b"N", b"SliceInt32", b"SliceInt64", b"SliceUInt32", b"SliceUInt64", b"SliceFloat32", b"SliceFloat64"])
231+
self.assertEqual(file["three/tree;1"].branchnames, [b"evt"])
232+
233+
self.assertEqual(dict((name, array.tolist()) for name, array in file["one/tree"].arrays(["one", "two", "three"]).items()), {b"one": [1, 2, 3, 4], b"two": [1.100000023841858, 2.200000047683716, 3.299999952316284, 4.400000095367432], b"three": [b"uno", b"dos", b"tres", b"quatro"]})
234+
self.assertEqual(file["one/two/tree"].array("Int32").shape, (100,))
235+
self.assertEqual(file["three/tree"].array("I32").shape, (100,))
236+
237+
file = uproot.open("tests/nesteddirs.root")
238+
239+
self.assertEqual(file["one/tree"].branchnames, [b"one", b"two", b"three"])
240+
self.assertEqual(file["one/two/tree"].branchnames, [b"Int32", b"Int64", b"UInt32", b"UInt64", b"Float32", b"Float64", b"Str", b"ArrayInt32", b"ArrayInt64", b"ArrayUInt32", b"ArrayUInt64", b"ArrayFloat32", b"ArrayFloat64", b"N", b"SliceInt32", b"SliceInt64", b"SliceUInt32", b"SliceUInt64", b"SliceFloat32", b"SliceFloat64"])
241+
self.assertEqual(file["three/tree"].branchnames, [b"evt"])
242+
243+
self.assertEqual(dict((name, array.tolist()) for name, array in file["one/tree;1"].arrays(["one", "two", "three"]).items()), {b"one": [1, 2, 3, 4], b"two": [1.100000023841858, 2.200000047683716, 3.299999952316284, 4.400000095367432], b"three": [b"uno", b"dos", b"tres", b"quatro"]})
244+
self.assertEqual(file["one/two/tree;1"].array("Int32").shape, (100,))
245+
self.assertEqual(file["three/tree;1"].array("I32").shape, (100,))

‎uproot/rootio.py

Lines changed: 35 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,9 @@ class TFile(object):
8080
"""Represents a ROOT file; use to extract objects.
8181
8282
* `file.get(name, cycle=None)` to extract an object (aware of '/' and ';' notations).
83-
* `file.ls()` for a `{name: classname}` dict of objects in the top directory.
84-
* `file.compression` for a Compression(algo, level) namedtuple.
83+
* `file.contents` is a `{name: classname}` dict of objects in the top directory.
84+
* `file.allcontents` is a `{name: classname}` dict of all objects in the file.
85+
* `file.compression` is a Compression(algo, level) namedtuple describing the compression.
8586
* `file.dir` is the top directory.
8687
8788
`file[name]` is a synonym for `file.get(name)`.
@@ -134,10 +135,13 @@ def __len__(self):
134135
def __iter__(self):
135136
return iter(self.dir.keys)
136137

137-
def ls(self):
138-
"""Get a `{name: classname}` dict of objects in the top directory.
139-
"""
140-
return self.dir.ls()
138+
@property
139+
def contents(self):
140+
return self.dir.contents
141+
142+
@property
143+
def allcontents(self):
144+
return self.dir.allcontents
141145

142146
def get(self, name, cycle=None):
143147
"""Get an object from the file, interpreting '/' as subdirectories and ';' to delimit cycle number.
@@ -152,7 +156,8 @@ class TDirectory(object):
152156
"""Represents a ROOT directory; use to extract objects.
153157
154158
* `dir.get(name, cycle=None)` to extract an object (aware of '/' and ';' notations).
155-
* `dir.ls()` for a `{name: classname}` dict of objects in this directory.
159+
* `dir.contents` is a `{name: classname}` dict of objects in this directory.
160+
* `dir.allcontents` is a `{name: classname}` dict of all objects under this directory.
156161
* `dir.keys` is the keys.
157162
158163
`dir[name]` is a synonym for `dir.get(name)`.
@@ -188,10 +193,13 @@ def __len__(self):
188193
def __iter__(self):
189194
return iter(self.keys)
190195

191-
def ls(self):
192-
"""Get a `{name: classname}` dict of objects in this directory.
193-
"""
194-
return self.keys.ls()
196+
@property
197+
def contents(self):
198+
return self.keys.contents
199+
200+
@property
201+
def allcontents(self):
202+
return self.keys.allcontents
195203

196204
def get(self, name, cycle=None):
197205
"""Get an object from the directory, interpreting '/' as subdirectories and ';' to delimit cycle number.
@@ -209,7 +217,8 @@ class TKeys(object):
209217
"""Represents a collection of keys.
210218
211219
* `keys.get(name, cycle=None)` to extract an object (aware of ';' notation).
212-
* `keys.ls()` for a `{name: classname}` dict.
220+
* `keys.contents` is a `{name: classname}` dict of objects directly in this set of TKeys.
221+
* `keys.allcontents` is a `{name: classname}` dict of all objects under this set of TKeys.
213222
214223
`keys[name]` is a synonym for `keys.get(name)`.
215224
@@ -240,10 +249,21 @@ def __len__(self):
240249
def __iter__(self):
241250
return iter(self.keys)
242251

243-
def ls(self):
244-
"""Get a `{name: classname}` dict.
252+
@property
253+
def contents(self):
254+
return dict(("{0};{1}".format(x.name.decode("ascii"), x.cycle).encode("ascii"), x.classname) for x in self.keys)
255+
256+
@property
257+
def allcontents(self):
258+
"""Get a `{name: classname}` dict of objects directly in this set of TKeys.
245259
"""
246-
return dict((x.name, x.classname) for x in self.keys)
260+
out = {}
261+
for name, classname in self.contents.items():
262+
out[name] = classname
263+
if classname == b"TDirectory":
264+
for name2, classname2 in self.get(name).allcontents.items():
265+
out["{0}/{1}".format(name[:name.rindex(b";")].decode("ascii"), name2.decode("ascii")).encode("ascii")] = classname2
266+
return out
247267

248268
def get(self, name, cycle=None):
249269
"""Get an object from the keys, interpreting ';' to delimit cycle number.

‎uproot/tree.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -247,13 +247,13 @@ def _normalizeselection(branchdtypes, allbranches):
247247
if name in allbranches:
248248
branch = allbranches[name]
249249
if hasattr(branch, "dtype"):
250-
yield branch, dtype
250+
yield branch, branch.dtype
251251
else:
252252
raise ValueError("cannot produce an array from branch {0}".format(repr(name)))
253253
else:
254254
raise ValueError("cannot find branch {0}".format(repr(name)))
255255

256-
def iterator(self, entries, branchdtypes=lambda branch: branch.dtype, executor=None, outputtype=dict, reportentries=False):
256+
def iterator(self, entries, branchdtypes=lambda branch: getattr(branch, "dtype", None), executor=None, outputtype=dict, reportentries=False):
257257
"""Iterates over a fixed number of entries at a time.
258258
259259
Instead of loading all entries from a tree with `tree.arrays()`, load a manageable number that will fit in memory at once and apply a continuous process to it. Example use:
@@ -344,7 +344,7 @@ def dobranch(branchdtypecache):
344344
else:
345345
yield out
346346

347-
def arrays(self, branchdtypes=lambda branch: branch.dtype, executor=None, outputtype=dict, block=True):
347+
def arrays(self, branchdtypes=lambda branch: getattr(branch, "dtype", None), executor=None, outputtype=dict, block=True):
348348
"""Extracts whole branches into Numpy arrays.
349349
350350
Individual branches from TTrees are typically small enough to fit into memory. If this is not your case, consider `tree.iterator(entries)` to load a given number of entries at a time.

0 commit comments

Comments
 (0)