Skip to content

Commit 5e6036c

Browse files
authored
Iterate over objects in TDirectory in linear time. (#638) (#639)
* Iterate over objects in TDirectory in linear time. * Remove the debug_counter. (cherry picked from commit 809cf8e)
1 parent 7d38054 commit 5e6036c

File tree

1 file changed

+22
-13
lines changed

1 file changed

+22
-13
lines changed

src/uproot/reading.py

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1444,12 +1444,19 @@ def __init__(self, path, cursor, context, file, parent):
14441444
)
14451445

14461446
self._keys = []
1447+
self._keys_lookup = {}
14471448
for _ in range(num_keys):
14481449
key = ReadOnlyKey(
14491450
keys_chunk, keys_cursor, {}, file, self, read_strings=True
14501451
)
1452+
name = key.fName
1453+
if name not in self._keys_lookup:
1454+
self._keys_lookup[name] = []
1455+
self._keys_lookup[name].append(len(self._keys))
14511456
self._keys.append(key)
14521457

1458+
self._len = None
1459+
14531460
self.hook_after_keys(
14541461
chunk=chunk,
14551462
cursor=cursor,
@@ -1901,11 +1908,13 @@ def _ipython_key_completions_(self):
19011908
return self.iterkeys() # noqa: B301 (not a dict)
19021909

19031910
def __len__(self):
1904-
return len(self._keys) + sum(
1905-
len(x.get())
1906-
for x in self._keys
1907-
if x.fClassName in ("TDirectory", "TDirectoryFile")
1908-
)
1911+
if self._len is None:
1912+
self._len = len(self._keys) + sum(
1913+
len(x.get())
1914+
for x in self._keys
1915+
if x.fClassName in ("TDirectory", "TDirectoryFile")
1916+
)
1917+
return self._len
19091918

19101919
def __contains__(self, where):
19111920
try:
@@ -2031,14 +2040,14 @@ def key(self, where):
20312040
item, cycle = where, None
20322041

20332042
last = None
2034-
for key in self._keys:
2035-
if key.fName == item:
2036-
if cycle == key.fCycle:
2037-
return key
2038-
elif cycle is None and last is None:
2039-
last = key
2040-
elif cycle is None and last.fCycle < key.fCycle:
2041-
last = key
2043+
for index in self._keys_lookup.get(item, []):
2044+
key = self._keys[index]
2045+
if cycle == key.fCycle:
2046+
return key
2047+
elif cycle is None and last is None:
2048+
last = key
2049+
elif cycle is None and last.fCycle < key.fCycle:
2050+
last = key
20422051

20432052
if last is not None:
20442053
return last

0 commit comments

Comments
 (0)