Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_ttree_empty_tbranch_multitree_uproot(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch("int32")
branchdict = {"intBranch": b}
tree = newtree(branchdict)
with uproot.recreate(filename, compression=None) as f:
for i in range(10):
f["t"*(i+1)] = tree
f = uproot.open(filename)
for i in range(10):
assert f["t" * (i + 1)]["intBranch"]._classname == b"TBranch"
def test_multi_branch_one_basket_diff_type(tmp_path):
filename = join(str(tmp_path), "example.root")
b1 = newbranch("int32")
b2 = newbranch("int64")
branchdict = {"intBranch": b1, "int8Branch": b2}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5], dtype=">i4")
b = numpy.array([6, 7, 8, 9, 10], dtype=">i8")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
f["t"]["intBranch"].newbasket(a)
f["t"]["int8Branch"].newbasket(b)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
intBranchdata = tree.AsMatrix(["intBranch"]).astype(">i4")
int8Branchdata = tree.AsMatrix(["int8Branch"]).astype(">i8")
for i in range(5):
assert a[i] == intBranchdata[i]
def test_multi_branch_one_basket_diff_type(tmp_path):
filename = join(str(tmp_path), "example.root")
b1 = newbranch("int32")
b2 = newbranch("int64")
branchdict = {"intBranch": b1, "int8Branch": b2}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5], dtype=">i4")
b = numpy.array([6, 7, 8, 9, 10], dtype=">i8")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
f["t"]["intBranch"].newbasket(a)
f["t"]["int8Branch"].newbasket(b)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
intBranchdata = tree.AsMatrix(["intBranch"]).astype(">i4")
int8Branchdata = tree.AsMatrix(["int8Branch"]).astype(">i8")
for i in range(5):
assert a[i] == intBranchdata[i]
assert b[i] == int8Branchdata[i]
def test_many_basket(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch(">i4")
branchdict = {"intBranch": b}
tree = newtree(branchdict)
a = numpy.array([1], dtype=">i4")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
for i in range(101):
f["t"]["intBranch"].newbasket(a)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
treedata = tree.AsMatrix().astype(">i4")
for i in range(101):
assert a[0] == treedata[i]
def test_many_basket_uproot(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch(">i4")
branchdict = {"intBranch": b}
tree = newtree(branchdict)
a = numpy.array([1], dtype=">i4")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
for i in range(101):
f["t"]["intBranch"].newbasket(a)
f = uproot.open(filename)
tree = f["t"]
treedata = tree.array("intBranch")
for i in range(101):
assert a[0] == treedata[i]
def test_multi_branch_one_basket_same_type(tmp_path):
filename = join(str(tmp_path), "example.root")
b1 = newbranch("int32")
b2 = newbranch("int32")
branchdict = {"intBranch": b1, "intBranch2": b2}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5], dtype=">i4")
b = numpy.array([6, 7, 8, 9, 10], dtype=">i4")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
f["t"]["intBranch"].newbasket(a)
f["t"]["intBranch2"].newbasket(b)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
intBranchdata = tree.AsMatrix(["intBranch"]).astype(">i4")
int8Branchdata = tree.AsMatrix(["intBranch2"]).astype(">i4")
for i in range(5):
assert a[i] == intBranchdata[i]
def test_tree_compression(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch(">i4")
branchdict = {"intBranch": b}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5], dtype=">i4")
with uproot.recreate(filename, compression=uproot.ZLIB(4)) as f:
f["t"] = tree
f["t"]["intBranch"].newbasket(a)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
treedata = tree.AsMatrix().astype(">i4")
for i in range(5):
assert a[i] == treedata[i]
def test_branch_basket_one_tleafs(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch(">i2")
branchdict = {"int2Branch": b}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5], dtype=">i2")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
f["t"]["int2Branch"].newbasket(a)
f = ROOT.TFile.Open(filename)
tree = f.Get("t")
treedata = tree.AsMatrix().astype(">i2")
for i in range(5):
assert a[i] == treedata[i]
def test_branch_basket_one_rewrite_root(tmp_path):
filename = join(str(tmp_path), "example.root")
b = newbranch("int32")
branchdict = {"intBranch": b}
tree = newtree(branchdict)
a = numpy.array([1, 2, 3, 4, 5]).astype("int32").newbyteorder(">")
with uproot.recreate(filename, compression=None) as f:
f["t"] = tree
f["t"]["intBranch"].newbasket(a)
f = ROOT.TFile.Open(filename, "UPDATE")
t = ROOT.TObjString("Hello World")
t.Write()
f.Close()
f = ROOT.TFile.Open(filename)
assert f.Get("Hello World") == "Hello World"
temp_arr[0:len(branch._branch.fields["_fBasketBytes"])] = branch._branch.fields["_fBasketBytes"]
branch._branch.fields["_fBasketBytes"] = temp_arr
tree = TTreeImpl(self._treelvl1._tree.name, newtree(), self._branch.file)
tree.name = self._treelvl1._tree.name
tree.fName = self._treelvl1._tree.fName
tree.fTitle = self._treelvl1._tree.fTitle
tree.fields["_fEntries"] = self._treelvl1._tree.fields["_fEntries"]
tree.fields["_fTotBytes"] = self._treelvl1._tree.fields["_fTotBytes"]
tree.fields["_fZipBytes"] = self._treelvl1._tree.fields["_fZipBytes"]
temp_branches = {}
for name, branch in self._treelvl1._branches.items():
compression = getattr(branch._branch, "compression", branch._branch.file.compression)
temp_branches[name] = TBranch(name, newbranch(branch._branch.type, ""), compression, self._treelvl1, branch._branch.file)
temp_branches[name]._branch.fields["_fWriteBasket"] = branch._branch.fields["_fWriteBasket"]
temp_branches[name]._branch.fields["_fEntries"] = branch._branch.fields["_fEntries"]
temp_branches[name]._branch.fields["_fBasketEntry"] = branch._branch.fields["_fBasketEntry"]
temp_branches[name]._branch.fields["_fEntryNumber"] = branch._branch.fields["_fEntryNumber"]
temp_branches[name]._branch.fields["_fMaxBaskets"] = branch._branch.fields["_fMaxBaskets"]
temp_branches[name]._branch.fields["_fBasketSeek"] = branch._branch.fields["_fBasketSeek"]
temp_branches[name]._branch.fields["_fBasketBytes"] = branch._branch.fields["_fBasketBytes"]
temp_branches[name]._branch.fields["_fTotBytes"] = branch._branch.fields["_fTotBytes"]
temp_branches[name]._branch.fields["_fZipBytes"] = branch._branch.fields["_fZipBytes"]
tree.fields["_fLeaves"].append(temp_branches[name]._branch.fields["_fLeaves"])
tree.fields["_fBranches"].append(temp_branches[name]._branch)
tree.branches = copy(temp_branches)
cursor = uproot.write.sink.cursor.Cursor(self._branch.file._fSeekFree)
tree.write_key = uproot.write.TKey.TKey(fClassName=self._treelvl1._tree.write_key.fClassName,