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

Commit

Permalink
Implemented TFileUpdate.bulk_write (#518)
Browse files Browse the repository at this point in the history
* Implemented TFileUpdate.bulk_write

* Renamed bulk_write to update and make it compatible to dict.update
  • Loading branch information
bfis committed Aug 28, 2020
1 parent 576a386 commit d9ea740
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 0 deletions.
19 changes: 19 additions & 0 deletions tests/test_write.py
Original file line number Diff line number Diff line change
Expand Up @@ -2134,3 +2134,22 @@ def test_jagged_i4_manybasket(tmp_path):
tree = f.Get("t")
for i, event in enumerate(tree):
assert(numpy.all([x for x in event.branch] == tester[i]))

def test_update(tmp_path):
filename = join(str(tmp_path), "example.root")
testfile = join(str(tmp_path), "test.root")
n = 3

f = ROOT.TFile.Open(testfile, "RECREATE")
h = ROOT.TH1F("hvar", "title", 5, 1, 10)
h.Write()
f.Close()

t = uproot.open(testfile)
hist = t["hvar"]
with uproot.recreate(filename, compression=None) as f:
f.update(("test%d" % i, hist) for i in range(n))

f = ROOT.TFile.Open(filename)
for i in range(n):
assert f.Get("test%d" % i).GetNbinsX() == 5, i
60 changes: 60 additions & 0 deletions uproot/write/TFile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@
import sys
import struct
import uuid
from itertools import chain
try:
from collections.abc import Mapping
except ImportError:
from collections import Mapping


import uproot_methods.convert

Expand Down Expand Up @@ -102,6 +108,60 @@ def __setitem__(self, where, what):
self._rootdir.setkey(newkey)
self._sink.flush()

def update(self, *args, **kwargs):
if len(args) > 1:
raise TypeError("update expected at most 1 argument, got %s" % len(args))
items = args[0] if args else ()
if isinstance(items, Mapping):
items = items.items()
items = chain(items, kwargs.items())
self.util = Util()

cursor = uproot.write.sink.cursor.Cursor(self._fSeekFree)
for where, what in items:
where, cycle = self._normalizewhere(where)

isTTree = what.__class__.__name__ in ("newtree", "TTree")
assert not isTTree # prevent TTree writing, otherwise migth invoke nasty magic
if not isTTree:
what = uproot_methods.convert.towriteable(what)
elif what.__class__.__name__ == "newtree":
what = TTree(where, what, self)

newkey = uproot.write.TKey.TKey(
fClassName=what._fClassName,
fName=where,
fTitle=what._fTitle,
fObjlen=0,
fSeekKey=cursor.index,
fSeekPdir=self._fBEGIN,
fCycle=cycle if cycle is not None else self._rootdir.newcycle(where),
)
if isTTree:
# Need to (re)attach the cycle number to allow getitem to access writable TTree
tree_where = where + b";" + str(newkey.fCycle).encode("utf-8")
self._treedict[tree_where] = what

newkeycursor = uproot.write.sink.cursor.Cursor(newkey.fSeekKey)
newkey.write(cursor, self._sink)
what._write(self, cursor, where, self.compression, newkey, newkeycursor, self.util)

dirkey = (newkey.fName, newkey.fCycle)
if dirkey in self._rootdir.keys:
self._rootdir.headkey.fObjlen -= self._rootdir.keys[dirkey].fKeylen
self._rootdir.headkey.fObjlen += newkey.fKeylen
self._rootdir.keys[dirkey] = newkey

# write (root) TDirectory
self._rootdir.fNbytesKeys = self._rootdir._nbyteskeys()
while self._rootdir.fNbytesKeys > self._rootdir.allocationbytes:
self._rootdir.allocationbytes *= self._rootdir.growfactor

self._rootdir.writekeys(cursor)

self._expandfile(cursor)
self._sink.flush()

def __delitem__(self, where):
where, cycle = self._normalizewhere(where)
try:
Expand Down

0 comments on commit d9ea740

Please sign in to comment.