Skip to content

Commit

Permalink
Fixed writing, all tests are passing so far
Browse files Browse the repository at this point in the history
  • Loading branch information
@zbilodea
zbilodea committed Jun 12, 2024
1 parent 7c9b96c commit 0b42253
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 79 deletions.
20 changes: 5 additions & 15 deletions src/uproot/writing/_cascadetree.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -997,6 +997,7 @@ def write_anew(self, sink):
absolute_location = key_num_bytes + sum(
len(x) for x in out if x is not None
)

absolute_location += 8 + 6 * (sum(1 if x is None else 0 for x in out) - 1)
datum["tleaf_reference_number"] = absolute_location + 2
tleaf_reference_numbers.append(datum["tleaf_reference_number"])
Expand Down Expand Up @@ -1206,7 +1207,6 @@ def write_anew(self, sink):
)
)
out.append(tleaf_reference_bytes)

# null fAliases (b"\x00\x00\x00\x00")
# empty fIndexValues array (4-byte length is zero)
# empty fIndex array (4-byte length is zero)
Expand Down Expand Up @@ -1685,7 +1685,6 @@ def write_with_new_branches(self, sink, old_key, directory):
# old_branches = uproot.writing._cascade.OldBranches(self._existing_branches)
for branch in self._existing_branches:
# # create OldTBranch object

cursor = (
branch.cursor.copy()
) # cursor before TObjArray of TBranches...hopefully
Expand All @@ -1704,16 +1703,6 @@ def write_with_new_branches(self, sink, old_key, directory):
+ 2
)

key_num_bytes = uproot.reading._key_format_big.size + 6
name_asbytes = branch.name.encode(errors="surrogateescape")
title_asbytes = branch.title.encode(errors="surrogateescape")
key_num_bytes += (1 if len(name_asbytes) < 255 else 5) + len(
name_asbytes
)
key_num_bytes += (1 if len(title_asbytes) < 255 else 5) + len(
title_asbytes
)

out.append(
self._existing_ttree.chunk.raw_data.tobytes()[
first_indx - branch_start : f_indx + 25
Expand All @@ -1726,7 +1715,9 @@ def write_with_new_branches(self, sink, old_key, directory):
absolute_location += 8 + 6 * (
sum(1 if x is None else 0 for x in out) - 1
)
tleaf_reference_numbers.append(absolute_location + 2)

tleaf_reference_numbers.append(absolute_location)

out.append(
self._existing_ttree.chunk.raw_data.tobytes()[
f_indx + 25 : second_indx
Expand Down Expand Up @@ -1957,7 +1948,6 @@ def write_with_new_branches(self, sink, old_key, directory):
datum["counter"]["tleaf_reference_number"]
)
)

# specialized TLeaf* members (fMinimum, fMaximum)
out.append(special_struct.pack(0, 0))

Expand Down Expand Up @@ -2019,7 +2009,7 @@ def write_with_new_branches(self, sink, old_key, directory):
tleaf_reference_bytes = uproot._util.tobytes(
numpy.array(tleaf_reference_numbers, ">u4")
)
out.append(
out.append( # This is still fine
struct.pack(
">I13sI4s",
(21 + len(tleaf_reference_bytes)) | uproot.const.kByteCountMask,
Expand Down
84 changes: 20 additions & 64 deletions tests/test_1155_feat_add_copy_ttree.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,14 +87,10 @@ def simple_test(tmp_path):
with uproot.open(
os.path.join(tmp_path, "arrays2.root"), minimal_ttree_metadata=False
) as new:
# print("???", new['whatever'].branches)
# print("???", check['whatever']['b1'].all_members)
# string = uproot.models.TString.Model_TString(check['whatever']['b1'].classname)
# print(type(string))
# print(new['whatever'].chunk.raw_data.tobytes())
# print(new["whatever"].chunk.raw_data.tobytes())
print(check["whatever"].chunk.raw_data.tobytes())
# print(check["whatever"].chunk.raw_data.tobytes())
for key in new["whatever"].keys():
assert ak.all(
new["whatever"].arrays()[key] == check["whatever"].arrays()[key]
)
assert ak.all(new["whatever"]["b1"].array() == data)
assert ak.all(new["whatever"]["b2"].array() == data1)
assert ak.all(new["whatever"]["b3"].array() == data)
Expand All @@ -103,6 +99,7 @@ def simple_test(tmp_path):
tree = inFile.Get("whatever;1")
indx = 0
for x in tree:
print(getattr(x, "b1"))
assert getattr(x, "b1") == data[indx]
assert getattr(x, "b2") == data1[indx]
indx += 1
Expand Down Expand Up @@ -173,7 +170,6 @@ def test_all_dtypes(tmp_path):


def test_ak_arrays(tmp_path):

data = np.array([1, 2, 3], dtype=np.int64)
data1 = np.array([2, 3, 4], dtype=np.int64)
data2 = np.array([3, 4, 5], dtype=np.int64)
Expand All @@ -187,13 +183,13 @@ def test_ak_arrays(tmp_path):
with uproot.recreate(os.path.join(tmp_path, "ak_test.root")) as file:
file["whatever"] = {
"b1": ak.Array([data, data1, data2]),
"b2": ak.Array([data1, data2, data]),
}

with uproot.update(os.path.join(tmp_path, "ak_test.root")) as write:
write.add_branches(
"whatever",
{
"b2": ak.Array([data1, data2, data]),
"b3": ak.Array([data2, data, data1]),
},
)
Expand All @@ -204,18 +200,20 @@ def test_ak_arrays(tmp_path):
with uproot.open(
os.path.join(tmp_path, "ak_test.root"), minimal_ttree_metadata=False
) as new:
print(correct["whatever"]["b3"].member("fLeaves")[0].all_members)
print(new["whatever"]["b3"].member("fLeaves")[0].all_members)
# inFile = ROOT.TFile.Open(os.path.join(tmp_path, "ak_test.root"), "READ")
# tree = inFile.Get("whatever")
# for x in tree:
# getattr(x, "b1")
# inFile.Close()
# df3 = ROOT.RDataFrame("whatever", os.path.join(tmp_path, "ak_test.root"))
# npy3 = ak.from_rdataframe(df3, columns=("b1", "b2", "b3"), keep_order=True)
# assert ak.all(npy3["b1"] == [data, data1, data2])
# assert ak.all(npy3["b2"] == [data1, data2, data])
# assert ak.all(npy3["b3"] == [data2, data, data1])
new["whatever"].arrays()
print(correct["whatever"]["b2"].member("fLeaves")[0].all_members)
print(new["whatever"]["b2"].member("fLeaves")[0].all_members)
print(new["whatever"]["b3"].arrays())
inFile = ROOT.TFile.Open(os.path.join(tmp_path, "ak_test.root"), "READ")
tree = inFile.Get("whatever")
for x in tree:
getattr(x, "b1")
inFile.Close()
df3 = ROOT.RDataFrame("whatever", os.path.join(tmp_path, "ak_test.root"))
npy3 = ak.from_rdataframe(df3, columns=("b1", "b2", "b3"), keep_order=True)
assert ak.all(npy3["b1"] == [data, data1, data2])
assert ak.all(npy3["b2"] == [data1, data2, data])
assert ak.all(npy3["b3"] == [data2, data, data1])


def test_streamers_same_dtypes(tmp_path):
Expand Down Expand Up @@ -341,42 +339,6 @@ def test_streamers_diff_dtypes(tmp_path):
inFile.Close()


def HZZ_test(tmp_path):
with uproot.open(
data_path("uproot-HZZ.root"), minimal_ttree_metadata=False
) as control:
with uproot.update(os.path.join(tmp_path, "uproot-HZZ.root copy")) as new:
data = []
for i in range(2421):
data.append(np.arange(0, 3, 1))
data = ak.Array(data)
new.add_branches("events", {"data": data})

with uproot.open(
os.path.join(tmp_path, "uproot-HZZ.root copy"),
minimal_ttree_metadata=False,
) as new:
for key in control["events"].keys():
assert key in new["events"].keys()
assert ak.all(
new["events"][key].array() == control["events"][key].array()
)
inFile = ROOT.TFile.Open(
os.path.join(tmp_path, "uproot-HZZ.root copy"), "READ"
)
tree = inFile.Get("events")
indx = 0
inFile.Close()
df3 = ROOT.RDataFrame(
"events", os.path.join(tmp_path, "uproot-HZZ.root copy")
)
npy3 = ak.from_rdataframe(df3, columns=("data"), keep_order=True)
# for key in npy3.keys():
# assert ak.all(npy3[key] == control['events'][key].array())
assert ak.all(npy3 == data)
inFile.Close()


def test_nested_branches(tmp_path):
# Make example
with uproot.open(data_path("uproot-HZZ-objects.root")):
Expand Down Expand Up @@ -424,9 +386,3 @@ def test_branch_v8(tmp_path):
with uproot.update(os.path.join(tmp_path, "uproot-issue-250.root")) as new:
print("hi")
# with uproot.open(os.path.join("uproot-from-geant4.root copy")) as new:


simple_test("/Users/zobil/Desktop/directory/")
# test_ak_arrays("/Users/zobil/Desktop/directory/vectorVector")

# Try uproot unit tests to generate uproot-events maybe?

0 comments on commit 0b42253

Please sign in to comment.