From 9e964666dd1776209603fb6748172d2c5f3a7109 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:16:50 +0200 Subject: [PATCH 01/38] Add in-place TTree branch addition --- src/uproot/writing/writable.py | 198 +++++++++++++++++++++++++++++++++ 1 file changed, 198 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index b99642cfe..7395235fe 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1599,6 +1599,204 @@ def copy_from( old_key.data_uncompressed_bytes, ) + def add_branches(self, source, branches): + """ + Args: + source (str): Name of existing TTree to add branches to. + branches (dict of str -> array): Names and data of new branches. + + Adds new branches to an existing TTree in-place. Only the new branch + data and an updated TTree header are written; existing data is never + touched. Works with both simple TBranch and TBranchElement files. + + .. code-block:: python + + with uproot.update("file.root") as f: + f.add_branches("tree", {"new_branch": np.ones(100, dtype=np.float32)}) + """ + import struct + import uproot.compression + + if self._file.sink.closed: + raise ValueError("cannot modify a TTree in a closed file") + + # open existing tree in read mode + existing_file = uproot.open(self.file_path, minimal_ttree_metadata=False) + try: + old_ttree = existing_file[source] + except Exception: + raise ValueError(f"TTree {source!r} not found in file {self.file_path}") from None + if not isinstance(old_ttree, uproot.TTree): + raise TypeError("'source' must be the name of a TTree") + + # get tree key info + tree_key = existing_file.key(source + ";1") + key_seek = tree_key.fSeekKey + key_len = tree_key.fKeylen + compression = existing_file._file.compression + file_end = existing_file._file.fEND + + # get directory key info + with uproot.update(self.file_path) as tmp: + dir_key = tmp._cascading.data.get_key(source, 1) + dir_key_location = dir_key.location + dir_key_big = dir_key.big + + # get decompressed blob + chunk, cursor = tree_key.get_uncompressed_chunk_cursor() + orig_raw = bytearray(chunk.raw_data.tobytes()) + num_branches = len(old_ttree.branches) + last_branch = list(old_ttree.branches)[-1] + c = last_branch.cursor.copy() + c.skip_after(last_branch) + insertion_point = c.index + existing_file.close() + + # find fBranches TObjArray bcnt + tobjarray_bcnt_pos = None + for i in range(190, 220): + val = struct.unpack(">I", orig_raw[i : i + 4])[0] + if val & 0x40000000 and (val & ~0x40000000) > 100: + tobjarray_bcnt_pos = i + old_tobjarray_bcnt = val & ~0x40000000 + break + if tobjarray_bcnt_pos is None: + raise RuntimeError("Could not find fBranches TObjArray byte count header") + + # build new blob inserting all new branches + new_blob = bytearray(orig_raw) + extra_bytes = 0 + + for branch_name, branch_data in branches.items(): + import numpy + + branch_data = numpy.asarray(branch_data) + dtype = branch_data.dtype + + # create minimal tree to get branch bytes + import tempfile, os + with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: + tmp_path = tmp_f.name + try: + with uproot.recreate(tmp_path) as tmp_file: + tmp_file.mktree("tree", {branch_name: dtype}) + tmp_file["tree"].extend({branch_name: branch_data}) + + with uproot.open(tmp_path) as tmp_open: + tmp_branch = tmp_open["tree"].branches[0] + basket_seek_val = tmp_branch.member("fBasketSeek")[0] + basket_bytes_size = tmp_branch.member("fBasketBytes")[0] + tmp_key = tmp_open.key("tree;1") + tmp_chunk, tmp_cursor = tmp_key.get_uncompressed_chunk_cursor() + tmp_raw = bytearray(tmp_chunk.raw_data.tobytes()) + tmp_fsize_pos = tmp_raw.find(struct.pack(">i", 1)) + tmp_c = tmp_branch.cursor.copy() + tmp_c.skip_after(tmp_branch) + tbranch_pos = tmp_raw.find(b"TBranch", tmp_fsize_pos) + elem_start = tbranch_pos - 8 + new_branch_bytes = bytearray(tmp_raw[elem_start : tmp_c.index]) + + # find fBasketSeek offset (8-byte) + target8 = struct.pack(">q", basket_seek_val) + idx8 = tmp_raw.find(target8, elem_start) + basket_seek_offset_8 = idx8 - elem_start + + # find tleaf offset + tleaf_fsize = tmp_raw.find(struct.pack(">i", 1), tmp_c.index) + tleaf_refs_start = tleaf_fsize + 8 + tleaf_ref = struct.unpack(">I", tmp_raw[tleaf_refs_start : tleaf_refs_start + 4])[0] + tleaf_offset = tleaf_ref - elem_start + + with open(tmp_path, "rb") as bf: + bf.seek(basket_seek_val) + basket_data_bytes = bytearray(bf.read(basket_bytes_size)) + finally: + os.unlink(tmp_path) + + # write basket at file_end, key after basket + new_basket_seek = file_end + new_key_seek = file_end + basket_bytes_size + + # update basket key header (8-byte fSeekKey) + struct.pack_into(">q", basket_data_bytes, 18, new_basket_seek) + + # insert new branch bytes at insertion point + new_blob = ( + new_blob[: insertion_point + extra_bytes] + + new_branch_bytes + + new_blob[insertion_point + extra_bytes :] + ) + + # patch fBasketSeek in blob + basket_seek_pos = insertion_point + extra_bytes + basket_seek_offset_8 + struct.pack_into(">q", new_blob, basket_seek_pos, new_basket_seek) + + # patch tLeaf fSize + tleaf_fsize_pos = new_blob.find( + struct.pack(">i", num_branches), insertion_point + extra_bytes + ) + struct.pack_into(">i", new_blob, tleaf_fsize_pos, num_branches + 1) + + # append tleaf ref + tleaf_refs_start_p = tleaf_fsize_pos + 8 + tleaf_refs_end = tleaf_refs_start_p + num_branches * 4 + new_tleaf_ref = struct.pack(">I", insertion_point + extra_bytes + tleaf_offset) + new_blob = new_blob[:tleaf_refs_end] + bytearray(new_tleaf_ref) + new_blob[tleaf_refs_end:] + + # patch tLeaf TObjArray bcnt + tleaf_tobjarray_bcnt_pos = insertion_point + extra_bytes + len(new_branch_bytes) + old_tleaf_bcnt = struct.unpack(">I", new_blob[tleaf_tobjarray_bcnt_pos : tleaf_tobjarray_bcnt_pos + 4])[0] & ~0x40000000 + struct.pack_into(">I", new_blob, tleaf_tobjarray_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000) + + extra_bytes += len(new_branch_bytes) + 4 # +4 for tleaf ref + + # write basket and update file_end for next branch + self._file.sink.write(new_basket_seek, bytes(basket_data_bytes)) + file_end = new_key_seek + + # patch TTree bcnt + total_added = extra_bytes + old_bcnt = struct.unpack(">I", new_blob[:4])[0] & ~0x40000000 + struct.pack_into(">I", new_blob, 0, (old_bcnt + total_added) | 0x40000000) + + # patch fBranches TObjArray bcnt + struct.pack_into(">I", new_blob, tobjarray_bcnt_pos, (old_tobjarray_bcnt + total_added - len(branches) * 4) | 0x40000000) + + # patch fBranches fSize + fbranches_fsize_pos = new_blob.find(struct.pack(">i", num_branches)) + struct.pack_into(">i", new_blob, fbranches_fsize_pos, num_branches + len(branches)) + + # compress and write new key + compressed = uproot.compression.compress(bytes(new_blob), compression) + new_nbytes = key_len + len(compressed) + new_objlen = len(new_blob) + + # copy original key header and update + self._file.sink.set_file_length(new_key_seek + new_nbytes) + raw_key = bytearray(self._file.sink.read(key_seek, key_len)) + struct.pack_into(">i", raw_key, 0, new_nbytes) + struct.pack_into(">i", raw_key, 6, new_objlen) + if dir_key_big: + struct.pack_into(">q", raw_key, 18, new_key_seek) + else: + struct.pack_into(">i", raw_key, 18, new_key_seek) + self._file.sink.write(new_key_seek, bytes(raw_key) + compressed) + + # update directory entry + raw_dir = bytearray(self._file.sink.read(dir_key_location, 26)) + struct.pack_into(">i", raw_dir, 0, new_nbytes) + struct.pack_into(">i", raw_dir, 6, new_objlen) + if dir_key_big: + struct.pack_into(">q", raw_dir, 18, new_key_seek) + else: + struct.pack_into(">i", raw_dir, 18, new_key_seek) + self._file.sink.write(dir_key_location, bytes(raw_dir)) + + # update fEND in file header + new_file_end = new_key_seek + new_nbytes + self._file.sink.write(12, struct.pack(">i", new_file_end)) + self._file.sink.flush() + def update(self, pairs=None, **more_pairs): """ Args: From 81ce9b3054f00c97fd6472d4caee2fb201a3f14c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:46:04 +0200 Subject: [PATCH 02/38] Add in-place TTree branch addition with tests --- src/uproot/writing/writable.py | 358 ++++++++++++++++----------------- tests/test_ttree_inplace.py | 128 ++++++++++++ 2 files changed, 305 insertions(+), 181 deletions(-) create mode 100644 tests/test_ttree_inplace.py diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 7395235fe..dafaa9df1 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1014,8 +1014,9 @@ def _get(self, name, cycle): if self._file._has_tree(key.seek_location): return self._file._get_tree(key.seek_location) else: - raise TypeError( - "WritableDirectory cannot view preexisting TTrees; open the file with uproot.open instead of uproot.recreate or uproot.update" + # return a WritableTree wrapper for preexisting trees (update mode) + return WritableTree( + self._path + (key.name.string,), self._file, None ) elif key.classname.string == "ROOT::RNTuple": if self._file._has_ntuple(key.seek_location): @@ -1599,37 +1600,165 @@ def copy_from( old_key.data_uncompressed_bytes, ) - def add_branches(self, source, branches): + def update(self, pairs=None, **more_pairs): + """ + Args: + pairs (dict or pairs of str \u2192 writable data): Names and data to write. + more_pairs (dict or pairs of str \u2192 writable data): More names and data to write. + + Bulk-update function, like assignment, but it collects TStreamerInfo for a single + update. + """ + streamers = [] + + if pairs is not None: + if hasattr(pairs, "keys"): + all_pairs = itertools.chain( + ((k, pairs[k]) for k in pairs.keys()), more_pairs.items() + ) + else: + all_pairs = itertools.chain(pairs, more_pairs.items()) + else: + all_pairs = more_pairs.items() + + for k, v in all_pairs: + fullpath = k.strip("/").split("/") + path, name = fullpath[:-1], fullpath[-1] + + if len(path) != 0: + self.mkdir( + "/".join(path), + initial_directory_bytes=self._file.initial_directory_bytes, + ) + + directory = self + for item in path: + directory = directory[item] + + uproot.writing.identify.add_to_directory(v, name, directory, streamers) + + self._file._cascading.streamers.update_streamers(self._file.sink, streamers) + + +class WritableTree: + """ + Args: + path (tuple of str): Path of directory names to this TTree. + file (:doc:`uproot.writing.writable.WritableFile`): Handle to the file in + which this TTree can be found. + cascading (:doc:`uproot.writing._cascadetree.Tree`): The low-level + directory object. + + Represents a writable ``TTree`` from a ROOT file. + + This object can be created using the :ref:`uproot.writing.writable.WritableDirectory.mktree` method. For instance: + + .. code-block:: python + + my_directory.mktree("tree1", {"branch1": np.array(...), "branch2": ak.Array(...)}) + my_directory.mktree("tree2", numpy_structured_array) + my_directory.mktree("tree3", awkward_record_array) + my_directory.mktree("tree4", pandas_dataframe) + + Recognized data types: + + * dict of NumPy arrays (flat, multidimensional, and/or structured), Awkward Arrays containing one level of variable-length lists and/or one level of records, or a Pandas DataFrame with a numeric index + * a single NumPy structured array (one level deep) + * a single Awkward Array containing one level of variable-length lists and/or one level of records + * a single Pandas DataFrame with a numeric index + + The arrays may have different types, but their lengths must be identical, at + least in the first dimension (i.e. number of entries). + + If the Awkward Array contains variable-length lists (i.e. it is "jagged"), a + counter TBranch will be created along with the data TBranch. ROOT needs the + counter TBranch to quantify the size of the variable-size arrays. Combining + Awkward Arrays with the same number of nested items using + `ak.zip `__ prevents + a proliferation of counter TBranches: + + .. code-block:: python + + my_directory.mktree("tree5", ak.zip({"branch1": array1, "branch2": array2, "branch3": array3})) + + would produce only one counter TBranch. + + The :doc:`uproot.writing.writable.WritableDirectory.mktree` method allows you to separate + the process of creating the TTree metadata from filling the first TBasket: + + .. code-block:: python + + my_directory.mktree("tree6", {"branch1": numpy_dtype, "branch2": awkward_type}) + + The :doc:`uproot.writing.writable.WritableDirectory.mktree` method can also control the + title of the TTree and the rules used to name counter TBranches and nested field TBranches. + + The ``numpy_dtype`` is any data that NumPy recognizes as a ``np.dtype``, and the + ``awkward_type`` is an `ak.types.Type `__ from + `ak.type `__ or + a string in that form, such as ``"var * float64"`` for variable-length doubles. + + TBaskets can be added to each TBranch using the :ref:`uproot.writing.writable.WritableTree.extend` + method: + + .. code-block:: python + + my_directory["tree6"].extend({"branch1": another_numpy_array, + "branch2": another_awkward_array}) + + Be sure to make these extensions as large as is feasible within memory constraints, + because a ROOT file full of small TBaskets is bloated (larger than it needs to be) + and slow to read (especially for Uproot, but also for ROOT). + + For instance, if you want to write a million events and have enough memory + available to do that 100 thousand events at a time (total of 10 TBaskets), + then do so. Filling the TTree a hundred events at a time (total of 10000 TBaskets) + would be considerably slower for writing and reading, and the file would be much + larger than it could otherwise be, even with compression. + """ + + def __init__(self, path, file, cascading): + self._path = path + self._file = file + self._cascading = cascading + + def add_branches(self, branches): """ Args: - source (str): Name of existing TTree to add branches to. branches (dict of str -> array): Names and data of new branches. - Adds new branches to an existing TTree in-place. Only the new branch - data and an updated TTree header are written; existing data is never - touched. Works with both simple TBranch and TBranchElement files. + Adds new branches to this TTree in-place. Only the new branch data and + an updated TTree header are written; existing data is never touched. + Works with both simple TBranch and TBranchElement files. .. code-block:: python with uproot.update("file.root") as f: - f.add_branches("tree", {"new_branch": np.ones(100, dtype=np.float32)}) + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) """ + import os import struct + import tempfile + + import numpy + import uproot.compression if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") + source = self._path[-1] + file_path = self._file.file_path + # open existing tree in read mode - existing_file = uproot.open(self.file_path, minimal_ttree_metadata=False) + existing_file = uproot.open(file_path, minimal_ttree_metadata=False) try: old_ttree = existing_file[source] except Exception: - raise ValueError(f"TTree {source!r} not found in file {self.file_path}") from None + raise ValueError(f"TTree {source!r} not found in file {file_path}") from None if not isinstance(old_ttree, uproot.TTree): raise TypeError("'source' must be the name of a TTree") - # get tree key info tree_key = existing_file.key(source + ";1") key_seek = tree_key.fSeekKey key_len = tree_key.fKeylen @@ -1637,12 +1766,11 @@ def add_branches(self, source, branches): file_end = existing_file._file.fEND # get directory key info - with uproot.update(self.file_path) as tmp: + with uproot.update(file_path) as tmp: dir_key = tmp._cascading.data.get_key(source, 1) dir_key_location = dir_key.location dir_key_big = dir_key.big - # get decompressed blob chunk, cursor = tree_key.get_uncompressed_chunk_cursor() orig_raw = bytearray(chunk.raw_data.tobytes()) num_branches = len(old_ttree.branches) @@ -1663,18 +1791,15 @@ def add_branches(self, source, branches): if tobjarray_bcnt_pos is None: raise RuntimeError("Could not find fBranches TObjArray byte count header") - # build new blob inserting all new branches new_blob = bytearray(orig_raw) extra_bytes = 0 + branch_extra_bytes = 0 + num_added = 0 for branch_name, branch_data in branches.items(): - import numpy - branch_data = numpy.asarray(branch_data) dtype = branch_data.dtype - # create minimal tree to get branch bytes - import tempfile, os with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: tmp_path = tmp_f.name try: @@ -1695,230 +1820,101 @@ def add_branches(self, source, branches): tbranch_pos = tmp_raw.find(b"TBranch", tmp_fsize_pos) elem_start = tbranch_pos - 8 new_branch_bytes = bytearray(tmp_raw[elem_start : tmp_c.index]) - - # find fBasketSeek offset (8-byte) target8 = struct.pack(">q", basket_seek_val) idx8 = tmp_raw.find(target8, elem_start) basket_seek_offset_8 = idx8 - elem_start - - # find tleaf offset tleaf_fsize = tmp_raw.find(struct.pack(">i", 1), tmp_c.index) tleaf_refs_start = tleaf_fsize + 8 tleaf_ref = struct.unpack(">I", tmp_raw[tleaf_refs_start : tleaf_refs_start + 4])[0] tleaf_offset = tleaf_ref - elem_start - with open(tmp_path, "rb") as bf: - bf.seek(basket_seek_val) - basket_data_bytes = bytearray(bf.read(basket_bytes_size)) + with open(tmp_path, "rb") as bf: + bf.seek(basket_seek_val) + basket_data_bytes = bytearray(bf.read(basket_bytes_size)) finally: os.unlink(tmp_path) - # write basket at file_end, key after basket new_basket_seek = file_end new_key_seek = file_end + basket_bytes_size # update basket key header (8-byte fSeekKey) struct.pack_into(">q", basket_data_bytes, 18, new_basket_seek) - # insert new branch bytes at insertion point - new_blob = ( - new_blob[: insertion_point + extra_bytes] - + new_branch_bytes - + new_blob[insertion_point + extra_bytes :] - ) + # insert new branch bytes + insert_at = insertion_point + branch_extra_bytes + new_blob = new_blob[:insert_at] + new_branch_bytes + new_blob[insert_at:] - # patch fBasketSeek in blob - basket_seek_pos = insertion_point + extra_bytes + basket_seek_offset_8 + # patch fBasketSeek + basket_seek_pos = insert_at + basket_seek_offset_8 struct.pack_into(">q", new_blob, basket_seek_pos, new_basket_seek) # patch tLeaf fSize + cur_num_branches = num_branches + num_added tleaf_fsize_pos = new_blob.find( - struct.pack(">i", num_branches), insertion_point + extra_bytes + struct.pack(">i", cur_num_branches), insert_at + len(new_branch_bytes) ) - struct.pack_into(">i", new_blob, tleaf_fsize_pos, num_branches + 1) + struct.pack_into(">i", new_blob, tleaf_fsize_pos, cur_num_branches + 1) # append tleaf ref tleaf_refs_start_p = tleaf_fsize_pos + 8 - tleaf_refs_end = tleaf_refs_start_p + num_branches * 4 - new_tleaf_ref = struct.pack(">I", insertion_point + extra_bytes + tleaf_offset) + tleaf_refs_end = tleaf_refs_start_p + cur_num_branches * 4 + new_tleaf_ref = struct.pack(">I", insert_at + tleaf_offset) new_blob = new_blob[:tleaf_refs_end] + bytearray(new_tleaf_ref) + new_blob[tleaf_refs_end:] # patch tLeaf TObjArray bcnt - tleaf_tobjarray_bcnt_pos = insertion_point + extra_bytes + len(new_branch_bytes) - old_tleaf_bcnt = struct.unpack(">I", new_blob[tleaf_tobjarray_bcnt_pos : tleaf_tobjarray_bcnt_pos + 4])[0] & ~0x40000000 - struct.pack_into(">I", new_blob, tleaf_tobjarray_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000) + tleaf_bcnt_pos = insert_at + len(new_branch_bytes) + old_tleaf_bcnt = struct.unpack(">I", new_blob[tleaf_bcnt_pos : tleaf_bcnt_pos + 4])[0] & ~0x40000000 + struct.pack_into(">I", new_blob, tleaf_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000) - extra_bytes += len(new_branch_bytes) + 4 # +4 for tleaf ref + branch_extra_bytes += len(new_branch_bytes) + extra_bytes += len(new_branch_bytes) + 4 + num_added += 1 - # write basket and update file_end for next branch + # write basket self._file.sink.write(new_basket_seek, bytes(basket_data_bytes)) file_end = new_key_seek # patch TTree bcnt - total_added = extra_bytes old_bcnt = struct.unpack(">I", new_blob[:4])[0] & ~0x40000000 - struct.pack_into(">I", new_blob, 0, (old_bcnt + total_added) | 0x40000000) + struct.pack_into(">I", new_blob, 0, (old_bcnt + extra_bytes) | 0x40000000) - # patch fBranches TObjArray bcnt - struct.pack_into(">I", new_blob, tobjarray_bcnt_pos, (old_tobjarray_bcnt + total_added - len(branches) * 4) | 0x40000000) + # patch fBranches TObjArray bcnt (extra_bytes minus tleaf refs) + struct.pack_into(">I", new_blob, tobjarray_bcnt_pos, + (old_tobjarray_bcnt + extra_bytes - num_added * 4) | 0x40000000) # patch fBranches fSize - fbranches_fsize_pos = new_blob.find(struct.pack(">i", num_branches)) - struct.pack_into(">i", new_blob, fbranches_fsize_pos, num_branches + len(branches)) + fbranches_fsize_pos = new_blob.find(struct.pack(">i", num_branches), tobjarray_bcnt_pos) + struct.pack_into(">i", new_blob, fbranches_fsize_pos, num_branches + num_added) # compress and write new key compressed = uproot.compression.compress(bytes(new_blob), compression) new_nbytes = key_len + len(compressed) new_objlen = len(new_blob) - # copy original key header and update - self._file.sink.set_file_length(new_key_seek + new_nbytes) raw_key = bytearray(self._file.sink.read(key_seek, key_len)) struct.pack_into(">i", raw_key, 0, new_nbytes) struct.pack_into(">i", raw_key, 6, new_objlen) if dir_key_big: - struct.pack_into(">q", raw_key, 18, new_key_seek) + struct.pack_into(">q", raw_key, 18, file_end) else: - struct.pack_into(">i", raw_key, 18, new_key_seek) - self._file.sink.write(new_key_seek, bytes(raw_key) + compressed) + struct.pack_into(">i", raw_key, 18, file_end) + self._file.sink.write(file_end, bytes(raw_key) + compressed) # update directory entry - raw_dir = bytearray(self._file.sink.read(dir_key_location, 26)) + raw_dir = bytearray(self._file.sink.read(dir_key_location, 40)) struct.pack_into(">i", raw_dir, 0, new_nbytes) struct.pack_into(">i", raw_dir, 6, new_objlen) if dir_key_big: - struct.pack_into(">q", raw_dir, 18, new_key_seek) + struct.pack_into(">q", raw_dir, 18, file_end) else: - struct.pack_into(">i", raw_dir, 18, new_key_seek) + struct.pack_into(">i", raw_dir, 18, file_end) self._file.sink.write(dir_key_location, bytes(raw_dir)) - # update fEND in file header - new_file_end = new_key_seek + new_nbytes + # update fEND + new_file_end = file_end + new_nbytes self._file.sink.write(12, struct.pack(">i", new_file_end)) self._file.sink.flush() - def update(self, pairs=None, **more_pairs): - """ - Args: - pairs (dict or pairs of str \u2192 writable data): Names and data to write. - more_pairs (dict or pairs of str \u2192 writable data): More names and data to write. - - Bulk-update function, like assignment, but it collects TStreamerInfo for a single - update. - """ - streamers = [] - - if pairs is not None: - if hasattr(pairs, "keys"): - all_pairs = itertools.chain( - ((k, pairs[k]) for k in pairs.keys()), more_pairs.items() - ) - else: - all_pairs = itertools.chain(pairs, more_pairs.items()) - else: - all_pairs = more_pairs.items() - - for k, v in all_pairs: - fullpath = k.strip("/").split("/") - path, name = fullpath[:-1], fullpath[-1] - - if len(path) != 0: - self.mkdir( - "/".join(path), - initial_directory_bytes=self._file.initial_directory_bytes, - ) - - directory = self - for item in path: - directory = directory[item] - - uproot.writing.identify.add_to_directory(v, name, directory, streamers) - - self._file._cascading.streamers.update_streamers(self._file.sink, streamers) - - -class WritableTree: - """ - Args: - path (tuple of str): Path of directory names to this TTree. - file (:doc:`uproot.writing.writable.WritableFile`): Handle to the file in - which this TTree can be found. - cascading (:doc:`uproot.writing._cascadetree.Tree`): The low-level - directory object. - - Represents a writable ``TTree`` from a ROOT file. - - This object can be created using the :ref:`uproot.writing.writable.WritableDirectory.mktree` method. For instance: - - .. code-block:: python - - my_directory.mktree("tree1", {"branch1": np.array(...), "branch2": ak.Array(...)}) - my_directory.mktree("tree2", numpy_structured_array) - my_directory.mktree("tree3", awkward_record_array) - my_directory.mktree("tree4", pandas_dataframe) - - Recognized data types: - - * dict of NumPy arrays (flat, multidimensional, and/or structured), Awkward Arrays containing one level of variable-length lists and/or one level of records, or a Pandas DataFrame with a numeric index - * a single NumPy structured array (one level deep) - * a single Awkward Array containing one level of variable-length lists and/or one level of records - * a single Pandas DataFrame with a numeric index - - The arrays may have different types, but their lengths must be identical, at - least in the first dimension (i.e. number of entries). - - If the Awkward Array contains variable-length lists (i.e. it is "jagged"), a - counter TBranch will be created along with the data TBranch. ROOT needs the - counter TBranch to quantify the size of the variable-size arrays. Combining - Awkward Arrays with the same number of nested items using - `ak.zip `__ prevents - a proliferation of counter TBranches: - - .. code-block:: python - - my_directory.mktree("tree5", ak.zip({"branch1": array1, "branch2": array2, "branch3": array3})) - - would produce only one counter TBranch. - - The :doc:`uproot.writing.writable.WritableDirectory.mktree` method allows you to separate - the process of creating the TTree metadata from filling the first TBasket: - - .. code-block:: python - - my_directory.mktree("tree6", {"branch1": numpy_dtype, "branch2": awkward_type}) - - The :doc:`uproot.writing.writable.WritableDirectory.mktree` method can also control the - title of the TTree and the rules used to name counter TBranches and nested field TBranches. - - The ``numpy_dtype`` is any data that NumPy recognizes as a ``np.dtype``, and the - ``awkward_type`` is an `ak.types.Type `__ from - `ak.type `__ or - a string in that form, such as ``"var * float64"`` for variable-length doubles. - - TBaskets can be added to each TBranch using the :ref:`uproot.writing.writable.WritableTree.extend` - method: - - .. code-block:: python - - my_directory["tree6"].extend({"branch1": another_numpy_array, - "branch2": another_awkward_array}) - - Be sure to make these extensions as large as is feasible within memory constraints, - because a ROOT file full of small TBaskets is bloated (larger than it needs to be) - and slow to read (especially for Uproot, but also for ROOT). - - For instance, if you want to write a million events and have enough memory - available to do that 100 thousand events at a time (total of 10 TBaskets), - then do so. Filling the TTree a hundred events at a time (total of 10000 TBaskets) - would be considerably slower for writing and reading, and the file would be much - larger than it could otherwise be, even with compression. - """ - - def __init__(self, path, file, cascading): - self._path = path - self._file = file - self._cascading = cascading - def __repr__(self): return "".format( repr("/" + "/".join(self._path)), id(self) diff --git a/tests/test_ttree_inplace.py b/tests/test_ttree_inplace.py new file mode 100644 index 000000000..52ecede8a --- /dev/null +++ b/tests/test_ttree_inplace.py @@ -0,0 +1,128 @@ +import os +import shutil + +import numpy as np +import pytest + +import uproot + +from skhep_testdata import data_path + +try: + import ROOT + + has_root = True +except ImportError: + has_root = False + +skip_no_root = pytest.mark.skipif(not has_root, reason="ROOT is not installed") + + +def test_add_branch_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 2 + assert "new_branch" in [b.name for b in f["tree"].branches] + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_multiple_branches(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({ + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + }) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 3 + assert np.all(f["tree"]["branch_a"].array() == 1.0) + assert np.all(f["tree"]["branch_b"].array() == 0) + + +def test_add_branch_int32(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_int": np.arange(100, dtype=np.int32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["new_int"].array() == np.arange(100, dtype=np.int32)) + + +def test_add_branch_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({ + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + }) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["x"].array() == np.arange(100, dtype=np.float32)) + assert np.all(f["tree"]["y"].array() == np.arange(100, dtype=np.int32)) + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_tbranchelement(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: + assert len(f["events"].branches) == 23 + assert np.all(f["events"]["new_branch"].array() == 1.0) + + +@skip_no_root +def test_add_branch_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree;1") + tree.SetCacheSize(0) + tree.GetEntry(0) + assert tree.new_branch == pytest.approx(1.0) + f.Close() + + +@skip_no_root +def test_add_branch_tbranchelement_root_readable(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") + tree = f.Get("events") + assert tree.GetNbranches() == 23 + branch = tree.GetBranch("new_branch") + assert branch.GetBasketSeek(0) > 0 + f.Close() From 3f89df640cea0bedbf39d60d4b28b9b1ebe0e28a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:47:58 +0200 Subject: [PATCH 03/38] Add in-place TTree branch addition with tests --- tests/test_ttree_inplace.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_ttree_inplace.py b/tests/test_ttree_inplace.py index 52ecede8a..0fcfddf11 100644 --- a/tests/test_ttree_inplace.py +++ b/tests/test_ttree_inplace.py @@ -123,6 +123,7 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") tree = f.Get("events") assert tree.GetNbranches() == 23 - branch = tree.GetBranch("new_branch") - assert branch.GetBasketSeek(0) > 0 + tree.SetCacheSize(0) + tree.GetEntry(0) + assert tree.new_branch == pytest.approx(1.0) f.Close() From 56fd418ebb089b8a070707d2c56d95c85ac0b1d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:33:30 +0200 Subject: [PATCH 04/38] Add in-place TTree extend method --- src/uproot/writing/writable.py | 161 +++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index dafaa9df1..83a406ea7 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1915,6 +1915,165 @@ def add_branches(self, branches): self._file.sink.write(12, struct.pack(">i", new_file_end)) self._file.sink.flush() + def _extend_inplace(self, data): + """ + Args: + data (dict of str -> array): Names and new data arrays for existing branches. + + Extends an existing TTree in-place by appending new entries to each branch. + Only new basket data and an updated TTree header are written; existing data + is never touched. Works with both simple TBranch and TBranchElement files. + + .. code-block:: python + + with uproot.update("file.root") as f: + f["tree"].extend({"x": np.ones(100, dtype=np.float32), + "y": np.zeros(100, dtype=np.int32)}) + """ + import os + import struct + import tempfile + + import numpy + + import uproot.compression + + if self._file.sink.closed: + raise ValueError("cannot modify a TTree in a closed file") + + source = self._path[-1] + file_path = self._file.file_path + + existing_file = uproot.open(file_path, minimal_ttree_metadata=False) + try: + old_ttree = existing_file[source] + except Exception: + raise ValueError(f"TTree {source!r} not found in file {file_path}") from None + + tree_key = existing_file.key(source + ";1") + key_seek = tree_key.fSeekKey + key_len = tree_key.fKeylen + compression = existing_file._file.compression + file_end = existing_file._file.fEND + + with uproot.update(file_path) as tmp: + dir_key = tmp._cascading.data.get_key(source, 1) + dir_key_location = dir_key.location + dir_key_big = dir_key.big + + chunk, cursor = tree_key.get_uncompressed_chunk_cursor() + orig_raw = bytearray(chunk.raw_data.tobytes()) + fEntries = old_ttree.member("fEntries") + fMaxBaskets = list(old_ttree.branches)[0].member("fMaxBaskets") + existing_file.close() + + # find TTree fEntries position in blob + fentries_pos = orig_raw.find(struct.pack(">q", fEntries)) + + # validate all branches exist and have same length + n_new = None + for bname, bdata in data.items(): + bdata = numpy.asarray(bdata) + if n_new is None: + n_new = len(bdata) + elif len(bdata) != n_new: + raise ValueError( + f"all arrays must have the same length, but {bname!r} has {len(bdata)} entries" + ) + + new_blob = bytearray(orig_raw) + current_file_end = file_end + + for bname, bdata in data.items(): + bdata = numpy.asarray(bdata) + + with uproot.open(file_path) as f: + branch = f[source][bname] + basket_seek_val = branch.member("fBasketSeek")[0] + fWriteBasket = branch.member("fWriteBasket") + + # find array positions from fBasketSeek[0] + target8 = struct.pack(">q", basket_seek_val) + seek_pos = new_blob.find(target8) + entry_pos = seek_pos - 1 - fMaxBaskets * 8 + bytes_pos = entry_pos - 1 - fMaxBaskets * 4 + + # find fWriteBasket position + wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack(">q", fEntries) + wb_pos = new_blob.find(wb_pattern, seek_pos - 500) + + # create new basket from temporary file + with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: + tmp_path = tmp_f.name + try: + with uproot.recreate(tmp_path) as tmp_file: + tmp_file.mktree("tree", {bname: bdata.dtype}) + tmp_file["tree"].extend({bname: bdata}) + with uproot.open(tmp_path) as tmp_open: + tmp_branch = tmp_open["tree"].branches[0] + new_basket_seek_val = tmp_branch.member("fBasketSeek")[0] + new_basket_bytes = tmp_branch.member("fBasketBytes")[0] + with open(tmp_path, "rb") as bf: + bf.seek(new_basket_seek_val) + basket_bytes_data = bytearray(bf.read(new_basket_bytes)) + finally: + os.unlink(tmp_path) + + new_basket_location = current_file_end + + # update basket key header fSeekKey (8-byte) + struct.pack_into(">q", basket_bytes_data, 18, new_basket_location) + + # patch blob + struct.pack_into(">i", new_blob, wb_pos, fWriteBasket + 1) # fWriteBasket + struct.pack_into(">q", new_blob, wb_pos + 4, fEntries + n_new) # fEntryNumber + struct.pack_into(">i", new_blob, bytes_pos + fWriteBasket * 4, new_basket_bytes) # fBasketBytes[fWriteBasket] + struct.pack_into(">q", new_blob, entry_pos + fWriteBasket * 8, fEntries) # fBasketEntry[fWriteBasket] + struct.pack_into(">q", new_blob, entry_pos + (fWriteBasket + 1) * 8, fEntries + n_new) # fBasketEntry[fWriteBasket+1] + struct.pack_into(">q", new_blob, seek_pos + fWriteBasket * 8, new_basket_location) # fBasketSeek[fWriteBasket] + + # patch branch fEntries (in _tbranch13_format2, after fSplitLevel=0) + branch_fentries_pattern = struct.pack(">i", 0) + struct.pack(">q", fEntries) + branch_fentries_pos = new_blob.find(branch_fentries_pattern, wb_pos) + 4 + struct.pack_into(">q", new_blob, branch_fentries_pos, fEntries + n_new) + + # write basket to file + self._file.sink.write(new_basket_location, bytes(basket_bytes_data)) + current_file_end = new_basket_location + new_basket_bytes + + # patch TTree fEntries + struct.pack_into(">q", new_blob, fentries_pos, fEntries + n_new) + + # compress and write new key + new_key_seek = current_file_end + compressed = uproot.compression.compress(bytes(new_blob), compression) + new_nbytes = key_len + len(compressed) + new_objlen = len(new_blob) + + raw_key = bytearray(self._file.sink.read(key_seek, key_len)) + struct.pack_into(">i", raw_key, 0, new_nbytes) + struct.pack_into(">i", raw_key, 6, new_objlen) + if dir_key_big: + struct.pack_into(">q", raw_key, 18, new_key_seek) + else: + struct.pack_into(">i", raw_key, 18, new_key_seek) + self._file.sink.write(new_key_seek, bytes(raw_key) + compressed) + + # update directory entry + raw_dir = bytearray(self._file.sink.read(dir_key_location, 40)) + struct.pack_into(">i", raw_dir, 0, new_nbytes) + struct.pack_into(">i", raw_dir, 6, new_objlen) + if dir_key_big: + struct.pack_into(">q", raw_dir, 18, new_key_seek) + else: + struct.pack_into(">i", raw_dir, 18, new_key_seek) + self._file.sink.write(dir_key_location, bytes(raw_dir)) + + # update fEND + new_file_end = new_key_seek + new_nbytes + self._file.sink.write(12, struct.pack(">i", new_file_end)) + self._file.sink.flush() + def __repr__(self): return "".format( repr("/" + "/".join(self._path)), id(self) @@ -2104,6 +2263,8 @@ def extend(self, data): **As a word of warning,** be sure that each call to :ref:`uproot.writing.writable.WritableTree.extend` includes at least 100 kB per branch/array. (NumPy and Awkward Arrays have an `nbytes `__ property; you want at least ``100000`` per array.) If you ask Uproot to write very small TBaskets, it will spend more time working on TBasket overhead than actually writing data. The absolute worst case is one-entry-per-:ref:`uproot.writing.writable.WritableTree.extend`. See `#428 (comment) `__. """ + if self._cascading is None: + return self._extend_inplace(data) self._cascading.extend(self._file, self._file.sink, data) def show( From 73d82ab9e497f5108e1a537861ca99d75941823b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:38:08 +0200 Subject: [PATCH 05/38] Add in-place TTree extend and add_branches with tests --- tests/test_ttree_inplace.py | 114 ++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/tests/test_ttree_inplace.py b/tests/test_ttree_inplace.py index 0fcfddf11..044dd3d67 100644 --- a/tests/test_ttree_inplace.py +++ b/tests/test_ttree_inplace.py @@ -127,3 +127,117 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): tree.GetEntry(0) assert tree.new_branch == pytest.approx(1.0) f.Close() + + +def test_extend_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["x"].array()[:100] == 1.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + assert np.all(f["tree"]["y"].array()[:100] == 0) + assert np.all(f["tree"]["y"].array()[100:] == 3) + + +def test_extend_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.arange(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.arange(100, dtype=np.float32) + 100}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + arr = f["tree"]["x"].array() + assert len(arr) == 200 + assert np.all(arr[:100] == np.arange(100, dtype=np.float32)) + assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) + + +def test_extend_tbranchelement(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) + + with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: + assert f["events"].member("fEntries") == 2521 + assert np.all(f["events"]["eventweight"].array()[2421:] == 99.0) + + +@skip_no_root +def test_extend_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree") + assert tree.GetEntries() == 150 + tree.SetCacheSize(0) + tree.GetEntry(149) + assert tree.x == pytest.approx(2.0) + f.Close() + + +@skip_no_root +def test_extend_tbranchelement_root_readable(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") + tree = f.Get("events") + assert tree.GetEntries() == 2521 + tree.SetCacheSize(0) + tree.GetEntry(2520) + assert tree.eventweight == pytest.approx(99.0) + f.Close() + + +def test_extend_nonexistent_branch(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) + + +def test_extend_mismatched_lengths(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="same length"): + f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) + + +def test_add_branch_nonexistent_tree(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) From 990acb0ee96771c14b28b904a9a453bbcca8a6d5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:42:41 +0000 Subject: [PATCH 06/38] style: pre-commit fixes --- src/uproot/writing/writable.py | 65 ++++++++++++++++++++++++---------- tests/test_ttree_inplace.py | 43 +++++++++++++++------- 2 files changed, 77 insertions(+), 31 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 83a406ea7..6794f964a 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1015,9 +1015,7 @@ def _get(self, name, cycle): return self._file._get_tree(key.seek_location) else: # return a WritableTree wrapper for preexisting trees (update mode) - return WritableTree( - self._path + (key.name.string,), self._file, None - ) + return WritableTree(self._path + (key.name.string,), self._file, None) elif key.classname.string == "ROOT::RNTuple": if self._file._has_ntuple(key.seek_location): return self._file._get_ntuple(key.seek_location) @@ -1755,7 +1753,9 @@ def add_branches(self, branches): try: old_ttree = existing_file[source] except Exception: - raise ValueError(f"TTree {source!r} not found in file {file_path}") from None + raise ValueError( + f"TTree {source!r} not found in file {file_path}" + ) from None if not isinstance(old_ttree, uproot.TTree): raise TypeError("'source' must be the name of a TTree") @@ -1825,7 +1825,9 @@ def add_branches(self, branches): basket_seek_offset_8 = idx8 - elem_start tleaf_fsize = tmp_raw.find(struct.pack(">i", 1), tmp_c.index) tleaf_refs_start = tleaf_fsize + 8 - tleaf_ref = struct.unpack(">I", tmp_raw[tleaf_refs_start : tleaf_refs_start + 4])[0] + tleaf_ref = struct.unpack( + ">I", tmp_raw[tleaf_refs_start : tleaf_refs_start + 4] + )[0] tleaf_offset = tleaf_ref - elem_start with open(tmp_path, "rb") as bf: @@ -1859,12 +1861,21 @@ def add_branches(self, branches): tleaf_refs_start_p = tleaf_fsize_pos + 8 tleaf_refs_end = tleaf_refs_start_p + cur_num_branches * 4 new_tleaf_ref = struct.pack(">I", insert_at + tleaf_offset) - new_blob = new_blob[:tleaf_refs_end] + bytearray(new_tleaf_ref) + new_blob[tleaf_refs_end:] + new_blob = ( + new_blob[:tleaf_refs_end] + + bytearray(new_tleaf_ref) + + new_blob[tleaf_refs_end:] + ) # patch tLeaf TObjArray bcnt tleaf_bcnt_pos = insert_at + len(new_branch_bytes) - old_tleaf_bcnt = struct.unpack(">I", new_blob[tleaf_bcnt_pos : tleaf_bcnt_pos + 4])[0] & ~0x40000000 - struct.pack_into(">I", new_blob, tleaf_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000) + old_tleaf_bcnt = ( + struct.unpack(">I", new_blob[tleaf_bcnt_pos : tleaf_bcnt_pos + 4])[0] + & ~0x40000000 + ) + struct.pack_into( + ">I", new_blob, tleaf_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000 + ) branch_extra_bytes += len(new_branch_bytes) extra_bytes += len(new_branch_bytes) + 4 @@ -1879,11 +1890,17 @@ def add_branches(self, branches): struct.pack_into(">I", new_blob, 0, (old_bcnt + extra_bytes) | 0x40000000) # patch fBranches TObjArray bcnt (extra_bytes minus tleaf refs) - struct.pack_into(">I", new_blob, tobjarray_bcnt_pos, - (old_tobjarray_bcnt + extra_bytes - num_added * 4) | 0x40000000) + struct.pack_into( + ">I", + new_blob, + tobjarray_bcnt_pos, + (old_tobjarray_bcnt + extra_bytes - num_added * 4) | 0x40000000, + ) # patch fBranches fSize - fbranches_fsize_pos = new_blob.find(struct.pack(">i", num_branches), tobjarray_bcnt_pos) + fbranches_fsize_pos = new_blob.find( + struct.pack(">i", num_branches), tobjarray_bcnt_pos + ) struct.pack_into(">i", new_blob, fbranches_fsize_pos, num_branches + num_added) # compress and write new key @@ -1948,7 +1965,9 @@ def _extend_inplace(self, data): try: old_ttree = existing_file[source] except Exception: - raise ValueError(f"TTree {source!r} not found in file {file_path}") from None + raise ValueError( + f"TTree {source!r} not found in file {file_path}" + ) from None tree_key = existing_file.key(source + ";1") key_seek = tree_key.fSeekKey @@ -2025,12 +2044,22 @@ def _extend_inplace(self, data): struct.pack_into(">q", basket_bytes_data, 18, new_basket_location) # patch blob - struct.pack_into(">i", new_blob, wb_pos, fWriteBasket + 1) # fWriteBasket - struct.pack_into(">q", new_blob, wb_pos + 4, fEntries + n_new) # fEntryNumber - struct.pack_into(">i", new_blob, bytes_pos + fWriteBasket * 4, new_basket_bytes) # fBasketBytes[fWriteBasket] - struct.pack_into(">q", new_blob, entry_pos + fWriteBasket * 8, fEntries) # fBasketEntry[fWriteBasket] - struct.pack_into(">q", new_blob, entry_pos + (fWriteBasket + 1) * 8, fEntries + n_new) # fBasketEntry[fWriteBasket+1] - struct.pack_into(">q", new_blob, seek_pos + fWriteBasket * 8, new_basket_location) # fBasketSeek[fWriteBasket] + struct.pack_into(">i", new_blob, wb_pos, fWriteBasket + 1) # fWriteBasket + struct.pack_into( + ">q", new_blob, wb_pos + 4, fEntries + n_new + ) # fEntryNumber + struct.pack_into( + ">i", new_blob, bytes_pos + fWriteBasket * 4, new_basket_bytes + ) # fBasketBytes[fWriteBasket] + struct.pack_into( + ">q", new_blob, entry_pos + fWriteBasket * 8, fEntries + ) # fBasketEntry[fWriteBasket] + struct.pack_into( + ">q", new_blob, entry_pos + (fWriteBasket + 1) * 8, fEntries + n_new + ) # fBasketEntry[fWriteBasket+1] + struct.pack_into( + ">q", new_blob, seek_pos + fWriteBasket * 8, new_basket_location + ) # fBasketSeek[fWriteBasket] # patch branch fEntries (in _tbranch13_format2, after fSplitLevel=0) branch_fentries_pattern = struct.pack(">i", 0) + struct.pack(">q", fEntries) diff --git a/tests/test_ttree_inplace.py b/tests/test_ttree_inplace.py index 044dd3d67..6b07275ed 100644 --- a/tests/test_ttree_inplace.py +++ b/tests/test_ttree_inplace.py @@ -38,10 +38,12 @@ def test_add_branch_multiple_branches(tmp_path): f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({ - "branch_a": np.ones(100, dtype=np.float32), - "branch_b": np.zeros(100, dtype=np.int32), - }) + f["tree"].add_branches( + { + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert len(f["tree"].branches) == 3 @@ -64,10 +66,12 @@ def test_add_branch_int32(tmp_path): def test_add_branch_preserves_existing(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({ - "x": np.arange(100, dtype=np.float32), - "y": np.arange(100, dtype=np.int32), - }) + f["tree"].extend( + { + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + } + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) @@ -132,10 +136,17 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): def test_extend_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "y": np.ones(50, dtype=np.int32) * 3, + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert f["tree"].member("fEntries") == 150 @@ -226,11 +237,15 @@ def test_extend_nonexistent_branch(tmp_path): def test_extend_mismatched_lengths(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="same length"): - f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} + ) def test_add_branch_nonexistent_tree(tmp_path): @@ -240,4 +255,6 @@ def test_add_branch_nonexistent_tree(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(Exception): - f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + f["nonexistent"].add_branches( + {"new_branch": np.ones(100, dtype=np.float32)} + ) From fbe47d898328c40253387b9bf5773ea37b2e1d07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:46:12 +0200 Subject: [PATCH 07/38] Rename test file to test_1690_ttree_inplace.py --- tests/test_1690_ttree_inplace.py | 243 +++++++++++++++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 tests/test_1690_ttree_inplace.py diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py new file mode 100644 index 000000000..044dd3d67 --- /dev/null +++ b/tests/test_1690_ttree_inplace.py @@ -0,0 +1,243 @@ +import os +import shutil + +import numpy as np +import pytest + +import uproot + +from skhep_testdata import data_path + +try: + import ROOT + + has_root = True +except ImportError: + has_root = False + +skip_no_root = pytest.mark.skipif(not has_root, reason="ROOT is not installed") + + +def test_add_branch_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 2 + assert "new_branch" in [b.name for b in f["tree"].branches] + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_multiple_branches(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({ + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + }) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 3 + assert np.all(f["tree"]["branch_a"].array() == 1.0) + assert np.all(f["tree"]["branch_b"].array() == 0) + + +def test_add_branch_int32(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_int": np.arange(100, dtype=np.int32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["new_int"].array() == np.arange(100, dtype=np.int32)) + + +def test_add_branch_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({ + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + }) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["x"].array() == np.arange(100, dtype=np.float32)) + assert np.all(f["tree"]["y"].array() == np.arange(100, dtype=np.int32)) + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_tbranchelement(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: + assert len(f["events"].branches) == 23 + assert np.all(f["events"]["new_branch"].array() == 1.0) + + +@skip_no_root +def test_add_branch_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree;1") + tree.SetCacheSize(0) + tree.GetEntry(0) + assert tree.new_branch == pytest.approx(1.0) + f.Close() + + +@skip_no_root +def test_add_branch_tbranchelement_root_readable(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") + tree = f.Get("events") + assert tree.GetNbranches() == 23 + tree.SetCacheSize(0) + tree.GetEntry(0) + assert tree.new_branch == pytest.approx(1.0) + f.Close() + + +def test_extend_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["x"].array()[:100] == 1.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + assert np.all(f["tree"]["y"].array()[:100] == 0) + assert np.all(f["tree"]["y"].array()[100:] == 3) + + +def test_extend_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.arange(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.arange(100, dtype=np.float32) + 100}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + arr = f["tree"]["x"].array() + assert len(arr) == 200 + assert np.all(arr[:100] == np.arange(100, dtype=np.float32)) + assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) + + +def test_extend_tbranchelement(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) + + with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: + assert f["events"].member("fEntries") == 2521 + assert np.all(f["events"]["eventweight"].array()[2421:] == 99.0) + + +@skip_no_root +def test_extend_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree") + assert tree.GetEntries() == 150 + tree.SetCacheSize(0) + tree.GetEntry(149) + assert tree.x == pytest.approx(2.0) + f.Close() + + +@skip_no_root +def test_extend_tbranchelement_root_readable(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), + os.path.join(tmp_path, "HZZ.root"), + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") + tree = f.Get("events") + assert tree.GetEntries() == 2521 + tree.SetCacheSize(0) + tree.GetEntry(2520) + assert tree.eventweight == pytest.approx(99.0) + f.Close() + + +def test_extend_nonexistent_branch(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) + + +def test_extend_mismatched_lengths(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="same length"): + f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) + + +def test_add_branch_nonexistent_tree(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) From cdf5d21b4a6199af7b8a90d4e8214d7ed1a6fef4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:47:19 +0000 Subject: [PATCH 08/38] style: pre-commit fixes --- tests/test_1690_ttree_inplace.py | 43 ++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 13 deletions(-) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 044dd3d67..6b07275ed 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -38,10 +38,12 @@ def test_add_branch_multiple_branches(tmp_path): f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({ - "branch_a": np.ones(100, dtype=np.float32), - "branch_b": np.zeros(100, dtype=np.int32), - }) + f["tree"].add_branches( + { + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert len(f["tree"].branches) == 3 @@ -64,10 +66,12 @@ def test_add_branch_int32(tmp_path): def test_add_branch_preserves_existing(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({ - "x": np.arange(100, dtype=np.float32), - "y": np.arange(100, dtype=np.int32), - }) + f["tree"].extend( + { + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + } + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) @@ -132,10 +136,17 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): def test_extend_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "y": np.ones(50, dtype=np.int32) * 3, + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert f["tree"].member("fEntries") == 150 @@ -226,11 +237,15 @@ def test_extend_nonexistent_branch(tmp_path): def test_extend_mismatched_lengths(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="same length"): - f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} + ) def test_add_branch_nonexistent_tree(tmp_path): @@ -240,4 +255,6 @@ def test_add_branch_nonexistent_tree(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(Exception): - f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + f["nonexistent"].add_branches( + {"new_branch": np.ones(100, dtype=np.float32)} + ) From daab666744e1376c293e36d96fca95e77364a60f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:49:53 +0200 Subject: [PATCH 09/38] Move imports to top level in writable.py --- src/uproot/writing/writable.py | 19 +-- tests/test_ttree_inplace.py | 260 --------------------------------- 2 files changed, 3 insertions(+), 276 deletions(-) delete mode 100644 tests/test_ttree_inplace.py diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 6794f964a..53d64bf81 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -20,8 +20,11 @@ import datetime import itertools +import os import queue +import struct import sys +import tempfile import uuid from collections.abc import Mapping, MutableMapping from pathlib import Path @@ -1734,14 +1737,6 @@ def add_branches(self, branches): with uproot.update("file.root") as f: f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) """ - import os - import struct - import tempfile - - import numpy - - import uproot.compression - if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") @@ -1947,14 +1942,6 @@ def _extend_inplace(self, data): f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) """ - import os - import struct - import tempfile - - import numpy - - import uproot.compression - if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") diff --git a/tests/test_ttree_inplace.py b/tests/test_ttree_inplace.py deleted file mode 100644 index 6b07275ed..000000000 --- a/tests/test_ttree_inplace.py +++ /dev/null @@ -1,260 +0,0 @@ -import os -import shutil - -import numpy as np -import pytest - -import uproot - -from skhep_testdata import data_path - -try: - import ROOT - - has_root = True -except ImportError: - has_root = False - -skip_no_root = pytest.mark.skipif(not has_root, reason="ROOT is not installed") - - -def test_add_branch_simple(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - assert len(f["tree"].branches) == 2 - assert "new_branch" in [b.name for b in f["tree"].branches] - assert np.all(f["tree"]["new_branch"].array() == 1.0) - - -def test_add_branch_multiple_branches(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches( - { - "branch_a": np.ones(100, dtype=np.float32), - "branch_b": np.zeros(100, dtype=np.int32), - } - ) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - assert len(f["tree"].branches) == 3 - assert np.all(f["tree"]["branch_a"].array() == 1.0) - assert np.all(f["tree"]["branch_b"].array() == 0) - - -def test_add_branch_int32(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({"new_int": np.arange(100, dtype=np.int32)}) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - assert np.all(f["tree"]["new_int"].array() == np.arange(100, dtype=np.int32)) - - -def test_add_branch_preserves_existing(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - { - "x": np.arange(100, dtype=np.float32), - "y": np.arange(100, dtype=np.int32), - } - ) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - assert np.all(f["tree"]["x"].array() == np.arange(100, dtype=np.float32)) - assert np.all(f["tree"]["y"].array() == np.arange(100, dtype=np.int32)) - assert np.all(f["tree"]["new_branch"].array() == 1.0) - - -def test_add_branch_tbranchelement(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) - - with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: - assert len(f["events"].branches) == 23 - assert np.all(f["events"]["new_branch"].array() == 1.0) - - -@skip_no_root -def test_add_branch_root_readable(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) - - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") - tree = f.Get("tree;1") - tree.SetCacheSize(0) - tree.GetEntry(0) - assert tree.new_branch == pytest.approx(1.0) - f.Close() - - -@skip_no_root -def test_add_branch_tbranchelement_root_readable(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") - tree = f.Get("events") - assert tree.GetNbranches() == 23 - tree.SetCacheSize(0) - tree.GetEntry(0) - assert tree.new_branch == pytest.approx(1.0) - f.Close() - - -def test_extend_simple(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} - ) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend( - { - "x": np.ones(50, dtype=np.float32) * 2, - "y": np.ones(50, dtype=np.int32) * 3, - } - ) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - assert f["tree"].member("fEntries") == 150 - assert np.all(f["tree"]["x"].array()[:100] == 1.0) - assert np.all(f["tree"]["x"].array()[100:] == 2.0) - assert np.all(f["tree"]["y"].array()[:100] == 0) - assert np.all(f["tree"]["y"].array()[100:] == 3) - - -def test_extend_preserves_existing(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.arange(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.arange(100, dtype=np.float32) + 100}) - - with uproot.open(os.path.join(tmp_path, "test.root")) as f: - arr = f["tree"]["x"].array() - assert len(arr) == 200 - assert np.all(arr[:100] == np.arange(100, dtype=np.float32)) - assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) - - -def test_extend_tbranchelement(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) - - with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: - assert f["events"].member("fEntries") == 2521 - assert np.all(f["events"]["eventweight"].array()[2421:] == 99.0) - - -@skip_no_root -def test_extend_root_readable(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") - tree = f.Get("tree") - assert tree.GetEntries() == 150 - tree.SetCacheSize(0) - tree.GetEntry(149) - assert tree.x == pytest.approx(2.0) - f.Close() - - -@skip_no_root -def test_extend_tbranchelement_root_readable(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") - tree = f.Get("events") - assert tree.GetEntries() == 2521 - tree.SetCacheSize(0) - tree.GetEntry(2520) - assert tree.eventweight == pytest.approx(99.0) - f.Close() - - -def test_extend_nonexistent_branch(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(Exception): - f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) - - -def test_extend_mismatched_lengths(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} - ) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(ValueError, match="same length"): - f["tree"].extend( - {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} - ) - - -def test_add_branch_nonexistent_tree(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(Exception): - f["nonexistent"].add_branches( - {"new_branch": np.ones(100, dtype=np.float32)} - ) From cf5532a445538b5507d052bb0ec834967148fcf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:05:35 +0200 Subject: [PATCH 10/38] Add accept_new_fields kwarg to extend --- src/uproot/writing/writable.py | 25 +++++++++++++++++++++---- tests/test_1690_ttree_inplace.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 53d64bf81..069b3eda7 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1927,7 +1927,7 @@ def add_branches(self, branches): self._file.sink.write(12, struct.pack(">i", new_file_end)) self._file.sink.flush() - def _extend_inplace(self, data): + def _extend_inplace(self, data, *, accept_new_fields=False): """ Args: data (dict of str -> array): Names and new data arrays for existing branches. @@ -1976,7 +1976,7 @@ def _extend_inplace(self, data): # find TTree fEntries position in blob fentries_pos = orig_raw.find(struct.pack(">q", fEntries)) - # validate all branches exist and have same length + # validate lengths and separate new vs existing branches n_new = None for bname, bdata in data.items(): bdata = numpy.asarray(bdata) @@ -1987,6 +1987,21 @@ def _extend_inplace(self, data): f"all arrays must have the same length, but {bname!r} has {len(bdata)} entries" ) + # handle new fields + existing_branch_names = [b.name for b in old_ttree.branches] + new_fields = {k: v for k, v in data.items() if k not in existing_branch_names} + if new_fields: + if not accept_new_fields: + raise ValueError( + f"new branches {list(new_fields.keys())} not in TTree; " + f"use accept_new_fields=True to add them automatically" + ) + # back-fill new branches with zeros for existing entries + zeros = {k: numpy.zeros(fEntries, dtype=numpy.asarray(v).dtype) for k, v in new_fields.items()} + self.add_branches(zeros) + # now extend all fields (existing + new) using fresh call + return self._extend_inplace(data, accept_new_fields=False) + new_blob = bytearray(orig_raw) current_file_end = file_end @@ -2252,10 +2267,12 @@ def num_baskets(self) -> int: """ return self._cascading.num_baskets - def extend(self, data): + def extend(self, data, *, accept_new_fields=False): """ Args: data (dict of str \u2192 arrays): More array data to add to the TTree. + accept_new_fields (bool): If True, new fields in data are automatically added + with zeros back-filled for existing entries before extending. This method adds data to an existing TTree, whether it was created through assignment or :doc:`uproot.writing.writable.WritableDirectory.mktree`. @@ -2280,7 +2297,7 @@ def extend(self, data): **As a word of warning,** be sure that each call to :ref:`uproot.writing.writable.WritableTree.extend` includes at least 100 kB per branch/array. (NumPy and Awkward Arrays have an `nbytes `__ property; you want at least ``100000`` per array.) If you ask Uproot to write very small TBaskets, it will spend more time working on TBasket overhead than actually writing data. The absolute worst case is one-entry-per-:ref:`uproot.writing.writable.WritableTree.extend`. See `#428 (comment) `__. """ if self._cascading is None: - return self._extend_inplace(data) + return self._extend_inplace(data, accept_new_fields=accept_new_fields) self._cascading.extend(self._file, self._file.sink, data) def show( diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 6b07275ed..2993ad22a 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -258,3 +258,32 @@ def test_add_branch_nonexistent_tree(tmp_path): f["nonexistent"].add_branches( {"new_branch": np.ones(100, dtype=np.float32)} ) + + +def test_extend_accept_new_fields(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend( + {"x": np.ones(50, dtype=np.float32) * 2, "new_branch": np.ones(50, dtype=np.float32) * 99}, + accept_new_fields=True, + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert "new_branch" in [b.name for b in f["tree"].branches] + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + + +def test_extend_new_fields_error_without_flag(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="accept_new_fields"): + f["tree"].extend({"x": np.ones(50, dtype=np.float32), "new_branch": np.ones(50, dtype=np.float32)}) From b0e7d4320aca357fb3f2d2f871b15bdad1e23059 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:06:03 +0000 Subject: [PATCH 11/38] style: pre-commit fixes --- src/uproot/writing/writable.py | 5 ++++- tests/test_1690_ttree_inplace.py | 12 ++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 069b3eda7..6a0e3b427 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1997,7 +1997,10 @@ def _extend_inplace(self, data, *, accept_new_fields=False): f"use accept_new_fields=True to add them automatically" ) # back-fill new branches with zeros for existing entries - zeros = {k: numpy.zeros(fEntries, dtype=numpy.asarray(v).dtype) for k, v in new_fields.items()} + zeros = { + k: numpy.zeros(fEntries, dtype=numpy.asarray(v).dtype) + for k, v in new_fields.items() + } self.add_branches(zeros) # now extend all fields (existing + new) using fresh call return self._extend_inplace(data, accept_new_fields=False) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 2993ad22a..82e7a4a8c 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -267,7 +267,10 @@ def test_extend_accept_new_fields(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].extend( - {"x": np.ones(50, dtype=np.float32) * 2, "new_branch": np.ones(50, dtype=np.float32) * 99}, + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + }, accept_new_fields=True, ) @@ -286,4 +289,9 @@ def test_extend_new_fields_error_without_flag(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="accept_new_fields"): - f["tree"].extend({"x": np.ones(50, dtype=np.float32), "new_branch": np.ones(50, dtype=np.float32)}) + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32), + "new_branch": np.ones(50, dtype=np.float32), + } + ) From 1b303df1228cb802712aeb07d9dc41ba807a1403 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:53:53 +0200 Subject: [PATCH 12/38] Fix extend for multiple sessions and add basket overflow check --- src/uproot/writing/writable.py | 31 ++++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 069b3eda7..fa498c165 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2012,6 +2012,7 @@ def _extend_inplace(self, data, *, accept_new_fields=False): branch = f[source][bname] basket_seek_val = branch.member("fBasketSeek")[0] fWriteBasket = branch.member("fWriteBasket") + fEntries_current = f[source].member("fEntries") # find array positions from fBasketSeek[0] target8 = struct.pack(">q", basket_seek_val) @@ -2019,9 +2020,33 @@ def _extend_inplace(self, data, *, accept_new_fields=False): entry_pos = seek_pos - 1 - fMaxBaskets * 8 bytes_pos = entry_pos - 1 - fMaxBaskets * 4 - # find fWriteBasket position - wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack(">q", fEntries) - wb_pos = new_blob.find(wb_pattern, seek_pos - 500) + # find fWriteBasket using fWriteBasket value read from file + # search for pattern: fWriteBasket(4) + fEntryNumber(8) near seek_pos + if fWriteBasket >= fMaxBaskets - 1: + raise ValueError( + f"branch {bname!r} has reached its maximum basket capacity ({fMaxBaskets}). " + f"Cannot extend further. Consider recreating the TTree with a larger initial_basket_capacity." + ) + + # read fEntries_current from new_blob (may have been updated in previous iteration) + fEntries_in_blob = struct.unpack(">q", new_blob[fentries_pos:fentries_pos+8])[0] + wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack(">q", fEntries_in_blob) + # search backward from seek_pos to find the LAST occurrence before seek_pos + wb_pos = -1 + search_start = max(0, seek_pos - 1000) + idx = search_start + while True: + idx = new_blob.find(wb_pattern, idx) + if idx == -1 or idx >= seek_pos: + break + wb_pos = idx + idx += 1 + if wb_pos == -1: + raise ValueError( + f"branch {bname!r} has likely reached its maximum basket capacity. " + f"Cannot extend further. Consider recreating the TTree with a larger initial_basket_capacity." + ) + entry_number_pos = wb_pos + 4 # create new basket from temporary file with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: From 314ee2fd16d70fadbc5ad784d265cc85cf656ed7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:57:32 +0000 Subject: [PATCH 13/38] style: pre-commit fixes --- src/uproot/writing/writable.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index f8a36c3d4..f03058f3b 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2032,8 +2032,12 @@ def _extend_inplace(self, data, *, accept_new_fields=False): ) # read fEntries_current from new_blob (may have been updated in previous iteration) - fEntries_in_blob = struct.unpack(">q", new_blob[fentries_pos:fentries_pos+8])[0] - wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack(">q", fEntries_in_blob) + fEntries_in_blob = struct.unpack( + ">q", new_blob[fentries_pos : fentries_pos + 8] + )[0] + wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack( + ">q", fEntries_in_blob + ) # search backward from seek_pos to find the LAST occurrence before seek_pos wb_pos = -1 search_start = max(0, seek_pos - 1000) From 2a2983d3b00efad69ab37f41f94ffbce172445af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:18:42 +0200 Subject: [PATCH 14/38] Use self._file instead of opening uproot.update again --- src/uproot/writing/writable.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index f8a36c3d4..1b9203b8a 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1760,11 +1760,10 @@ def add_branches(self, branches): compression = existing_file._file.compression file_end = existing_file._file.fEND - # get directory key info - with uproot.update(file_path) as tmp: - dir_key = tmp._cascading.data.get_key(source, 1) - dir_key_location = dir_key.location - dir_key_big = dir_key.big + # get directory key info from current file + dir_key = self._file._cascading.rootdirectory.data.get_key(source, 1) + dir_key_location = dir_key.location + dir_key_big = dir_key.big chunk, cursor = tree_key.get_uncompressed_chunk_cursor() orig_raw = bytearray(chunk.raw_data.tobytes()) @@ -1962,10 +1961,10 @@ def _extend_inplace(self, data, *, accept_new_fields=False): compression = existing_file._file.compression file_end = existing_file._file.fEND - with uproot.update(file_path) as tmp: - dir_key = tmp._cascading.data.get_key(source, 1) - dir_key_location = dir_key.location - dir_key_big = dir_key.big + # get directory key info from current file + dir_key = self._file._cascading.rootdirectory.data.get_key(source, 1) + dir_key_location = dir_key.location + dir_key_big = dir_key.big chunk, cursor = tree_key.get_uncompressed_chunk_cursor() orig_raw = bytearray(chunk.raw_data.tobytes()) From cbb47aa658196393990bf72919c9982439a66a63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:25:08 +0200 Subject: [PATCH 15/38] Fix hardcoded byte range for fBranches TObjArray bcnt search --- src/uproot/writing/writable.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index c08952767..4fe882bf9 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1776,12 +1776,16 @@ def add_branches(self, branches): # find fBranches TObjArray bcnt tobjarray_bcnt_pos = None - for i in range(190, 220): + # TObjArray bcnt: 4-byte value with 0x40000000 (kByteCountMask) bit set, + # immediately followed by 2-byte version=3 + for i in range(len(orig_raw) - 4 - 2): # -4 for bcnt, -2 for version val = struct.unpack(">I", orig_raw[i : i + 4])[0] if val & 0x40000000 and (val & ~0x40000000) > 100: - tobjarray_bcnt_pos = i - old_tobjarray_bcnt = val & ~0x40000000 - break + version = struct.unpack(">H", orig_raw[i + 4 : i + 6])[0] + if version == 3: + tobjarray_bcnt_pos = i + old_tobjarray_bcnt = val & ~0x40000000 + break if tobjarray_bcnt_pos is None: raise RuntimeError("Could not find fBranches TObjArray byte count header") From 36429fea433d0dfa4840b79a2df59e0c80ce5ec2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:41:24 +0200 Subject: [PATCH 16/38] Find TTree fEntries more reliably using unique sequence --- src/uproot/writing/writable.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 4fe882bf9..df256a57e 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1977,8 +1977,12 @@ def _extend_inplace(self, data, *, accept_new_fields=False): existing_file.close() # find TTree fEntries position in blob - fentries_pos = orig_raw.find(struct.pack(">q", fEntries)) - + fTotBytes = old_ttree.member("fTotBytes") + fZipBytes = old_ttree.member("fZipBytes") + fentries_seq = struct.pack(">q", fEntries) + struct.pack(">q", fTotBytes) + struct.pack(">q", fZipBytes) + fentries_pos = orig_raw.find(fentries_seq) + if fentries_pos == -1: + raise RuntimeError("Could not find TTree fEntries position in blob") # validate lengths and separate new vs existing branches n_new = None for bname, bdata in data.items(): From ba5411468d74615f95b4dffad862308f35e015e0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:42:09 +0000 Subject: [PATCH 17/38] style: pre-commit fixes --- src/uproot/writing/writable.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index df256a57e..12534d1f1 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1979,7 +1979,11 @@ def _extend_inplace(self, data, *, accept_new_fields=False): # find TTree fEntries position in blob fTotBytes = old_ttree.member("fTotBytes") fZipBytes = old_ttree.member("fZipBytes") - fentries_seq = struct.pack(">q", fEntries) + struct.pack(">q", fTotBytes) + struct.pack(">q", fZipBytes) + fentries_seq = ( + struct.pack(">q", fEntries) + + struct.pack(">q", fTotBytes) + + struct.pack(">q", fZipBytes) + ) fentries_pos = orig_raw.find(fentries_seq) if fentries_pos == -1: raise RuntimeError("Could not find TTree fEntries position in blob") From f012ece7e1dff3697b0c477159440b82f1d23f52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:45:55 +0200 Subject: [PATCH 18/38] Validate branch length matches tree in add_branches --- src/uproot/writing/writable.py | 9 ++++++++- tests/test_1690_ttree_inplace.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index df256a57e..cbb36bec9 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1772,8 +1772,15 @@ def add_branches(self, branches): c = last_branch.cursor.copy() c.skip_after(last_branch) insertion_point = c.index + tree_entries = old_ttree.member("fEntries") existing_file.close() - + # validate all new branches have same length as existing tree + for bname, bdata in branches.items(): + if len(numpy.asarray(bdata)) != tree_entries: + raise ValueError( + f"branch {bname!r} has {len(numpy.asarray(bdata))} entries but TTree has " + f"{tree_entries} entries; all new branches must match the tree length" + ) # find fBranches TObjArray bcnt tobjarray_bcnt_pos = None # TObjArray bcnt: 4-byte value with 0x40000000 (kByteCountMask) bit set, diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 82e7a4a8c..dc4086398 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -295,3 +295,13 @@ def test_extend_new_fields_error_without_flag(tmp_path): "new_branch": np.ones(50, dtype=np.float32), } ) + + +def test_add_branch_wrong_length(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="entries"): + f["tree"].add_branches({"new_branch": np.ones(50, dtype=np.float32)}) From 63732155ad22f003d51b85948f6af278a54c7e38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:09:51 +0200 Subject: [PATCH 19/38] Clean up tests and enforce all branches in extend --- src/uproot/writing/writable.py | 6 ++ tests/test_1690_ttree_inplace.py | 171 +++++++++++-------------------- 2 files changed, 65 insertions(+), 112 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d738540ff..24de142f3 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2008,6 +2008,12 @@ def _extend_inplace(self, data, *, accept_new_fields=False): # handle new fields existing_branch_names = [b.name for b in old_ttree.branches] new_fields = {k: v for k, v in data.items() if k not in existing_branch_names} + # check all existing branches are present (partial extends are inconsistent) + missing = [b for b in existing_branch_names if b not in data] + if missing: + raise ValueError( + f"data is missing branches {missing}; all existing branches must be extended together" + ) if new_fields: if not accept_new_fields: raise ValueError( diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index dc4086398..f58aeb3dc 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -5,12 +5,10 @@ import pytest import uproot - from skhep_testdata import data_path try: import ROOT - has_root = True except ImportError: has_root = False @@ -18,6 +16,8 @@ skip_no_root = pytest.mark.skipif(not has_root, reason="ROOT is not installed") +# ── add_branches tests ──────────────────────────────────────────────────────── + def test_add_branch_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32}) @@ -38,12 +38,10 @@ def test_add_branch_multiple_branches(tmp_path): f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches( - { - "branch_a": np.ones(100, dtype=np.float32), - "branch_b": np.zeros(100, dtype=np.int32), - } - ) + f["tree"].add_branches({ + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + }) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert len(f["tree"].branches) == 3 @@ -66,12 +64,10 @@ def test_add_branch_int32(tmp_path): def test_add_branch_preserves_existing(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - { - "x": np.arange(100, dtype=np.float32), - "y": np.arange(100, dtype=np.int32), - } - ) + f["tree"].extend({ + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + }) with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) @@ -83,10 +79,7 @@ def test_add_branch_preserves_existing(tmp_path): def test_add_branch_tbranchelement(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) + shutil.copy(data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root")) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) @@ -96,6 +89,26 @@ def test_add_branch_tbranchelement(tmp_path): assert np.all(f["events"]["new_branch"].array() == 1.0) +def test_add_branch_wrong_length(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="entries"): + f["tree"].add_branches({"new_branch": np.ones(50, dtype=np.float32)}) + + +def test_add_branch_nonexistent_tree(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + @skip_no_root def test_add_branch_root_readable(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: @@ -115,10 +128,7 @@ def test_add_branch_root_readable(tmp_path): @skip_no_root def test_add_branch_tbranchelement_root_readable(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) + shutil.copy(data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root")) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) @@ -133,20 +143,15 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): f.Close() +# ── extend tests ────────────────────────────────────────────────────────────── + def test_extend_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} - ) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend( - { - "x": np.ones(50, dtype=np.float32) * 2, - "y": np.ones(50, dtype=np.int32) * 3, - } - ) + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert f["tree"].member("fEntries") == 150 @@ -171,93 +176,34 @@ def test_extend_preserves_existing(tmp_path): assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) -def test_extend_tbranchelement(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) - - with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: - assert f["events"].member("fEntries") == 2521 - assert np.all(f["events"]["eventweight"].array()[2421:] == 99.0) - - -@skip_no_root -def test_extend_root_readable(tmp_path): +def test_extend_missing_branch(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) - - with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") - tree = f.Get("tree") - assert tree.GetEntries() == 150 - tree.SetCacheSize(0) - tree.GetEntry(149) - assert tree.x == pytest.approx(2.0) - f.Close() - - -@skip_no_root -def test_extend_tbranchelement_root_readable(tmp_path): - shutil.copy( - data_path("uproot-HZZ-objects.root"), - os.path.join(tmp_path, "HZZ.root"), - ) - - with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].extend({"eventweight": np.ones(100, dtype=np.float32) * 99.0}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") - tree = f.Get("events") - assert tree.GetEntries() == 2521 - tree.SetCacheSize(0) - tree.GetEntry(2520) - assert tree.eventweight == pytest.approx(99.0) - f.Close() - - -def test_extend_nonexistent_branch(tmp_path): - with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: - f.mktree("tree", {"x": np.float32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(Exception): - f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) + with pytest.raises(ValueError, match="missing"): + f["tree"].extend({"x": np.ones(50, dtype=np.float32)}) def test_extend_mismatched_lengths(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend( - {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} - ) + f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="same length"): - f["tree"].extend( - {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} - ) + f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) -def test_add_branch_nonexistent_tree(tmp_path): +def test_extend_nonexistent_branch(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32}) f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(Exception): - f["nonexistent"].add_branches( - {"new_branch": np.ones(100, dtype=np.float32)} - ) + f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) def test_extend_accept_new_fields(tmp_path): @@ -267,10 +213,7 @@ def test_extend_accept_new_fields(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].extend( - { - "x": np.ones(50, dtype=np.float32) * 2, - "new_branch": np.ones(50, dtype=np.float32) * 99, - }, + {"x": np.ones(50, dtype=np.float32) * 2, "new_branch": np.ones(50, dtype=np.float32) * 99}, accept_new_fields=True, ) @@ -289,19 +232,23 @@ def test_extend_new_fields_error_without_flag(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="accept_new_fields"): - f["tree"].extend( - { - "x": np.ones(50, dtype=np.float32), - "new_branch": np.ones(50, dtype=np.float32), - } - ) + f["tree"].extend({"x": np.ones(50, dtype=np.float32), "new_branch": np.ones(50, dtype=np.float32)}) -def test_add_branch_wrong_length(tmp_path): +@skip_no_root +def test_extend_root_readable(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32}) f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(ValueError, match="entries"): - f["tree"].add_branches({"new_branch": np.ones(50, dtype=np.float32)}) + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree") + assert tree.GetEntries() == 150 + tree.SetCacheSize(0) + tree.GetEntry(149) + assert tree.x == pytest.approx(2.0) + f.Close() \ No newline at end of file From f97c03ddeb69a6df9970ef316a1ad8b63bf4a289 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:10:18 +0000 Subject: [PATCH 20/38] style: pre-commit fixes --- tests/test_1690_ttree_inplace.py | 72 +++++++++++++++++++++++--------- 1 file changed, 53 insertions(+), 19 deletions(-) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index f58aeb3dc..a5c3f5a7a 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -9,6 +9,7 @@ try: import ROOT + has_root = True except ImportError: has_root = False @@ -18,6 +19,7 @@ # ── add_branches tests ──────────────────────────────────────────────────────── + def test_add_branch_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32}) @@ -38,10 +40,12 @@ def test_add_branch_multiple_branches(tmp_path): f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].add_branches({ - "branch_a": np.ones(100, dtype=np.float32), - "branch_b": np.zeros(100, dtype=np.int32), - }) + f["tree"].add_branches( + { + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert len(f["tree"].branches) == 3 @@ -64,10 +68,12 @@ def test_add_branch_int32(tmp_path): def test_add_branch_preserves_existing(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({ - "x": np.arange(100, dtype=np.float32), - "y": np.arange(100, dtype=np.int32), - }) + f["tree"].extend( + { + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + } + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) @@ -79,7 +85,9 @@ def test_add_branch_preserves_existing(tmp_path): def test_add_branch_tbranchelement(tmp_path): - shutil.copy(data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root")) + shutil.copy( + data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") + ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) @@ -106,7 +114,9 @@ def test_add_branch_nonexistent_tree(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(Exception): - f["nonexistent"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + f["nonexistent"].add_branches( + {"new_branch": np.ones(100, dtype=np.float32)} + ) @skip_no_root @@ -128,7 +138,9 @@ def test_add_branch_root_readable(tmp_path): @skip_no_root def test_add_branch_tbranchelement_root_readable(tmp_path): - shutil.copy(data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root")) + shutil.copy( + data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") + ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) @@ -145,13 +157,21 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): # ── extend tests ────────────────────────────────────────────────────────────── + def test_extend_simple(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2, "y": np.ones(50, dtype=np.int32) * 3}) + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "y": np.ones(50, dtype=np.int32) * 3, + } + ) with uproot.open(os.path.join(tmp_path, "test.root")) as f: assert f["tree"].member("fEntries") == 150 @@ -179,7 +199,9 @@ def test_extend_preserves_existing(tmp_path): def test_extend_missing_branch(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="missing"): @@ -189,11 +211,15 @@ def test_extend_missing_branch(tmp_path): def test_extend_mismatched_lengths(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: f.mktree("tree", {"x": np.float32, "y": np.int32}) - f["tree"].extend({"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="same length"): - f["tree"].extend({"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)}) + f["tree"].extend( + {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} + ) def test_extend_nonexistent_branch(tmp_path): @@ -213,7 +239,10 @@ def test_extend_accept_new_fields(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: f["tree"].extend( - {"x": np.ones(50, dtype=np.float32) * 2, "new_branch": np.ones(50, dtype=np.float32) * 99}, + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + }, accept_new_fields=True, ) @@ -232,7 +261,12 @@ def test_extend_new_fields_error_without_flag(tmp_path): with uproot.update(os.path.join(tmp_path, "test.root")) as f: with pytest.raises(ValueError, match="accept_new_fields"): - f["tree"].extend({"x": np.ones(50, dtype=np.float32), "new_branch": np.ones(50, dtype=np.float32)}) + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32), + "new_branch": np.ones(50, dtype=np.float32), + } + ) @skip_no_root @@ -251,4 +285,4 @@ def test_extend_root_readable(tmp_path): tree.SetCacheSize(0) tree.GetEntry(149) assert tree.x == pytest.approx(2.0) - f.Close() \ No newline at end of file + f.Close() From 314fe057869063db06123d364d1a231cccef7731 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:13:17 +0200 Subject: [PATCH 21/38] Fix fEND write size for big files (>2GB) --- src/uproot/writing/writable.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 24de142f3..e27d28bbd 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1934,7 +1934,11 @@ def add_branches(self, branches): # update fEND new_file_end = file_end + new_nbytes - self._file.sink.write(12, struct.pack(">i", new_file_end)) + # fEND is 4-byte for small files, 8-byte for files >= 2GB + if self._file._cascading.fileheader.big: + self._file.sink.write(12, struct.pack(">q", new_file_end)) + else: + self._file.sink.write(12, struct.pack(">i", new_file_end)) self._file.sink.flush() def _extend_inplace(self, data, *, accept_new_fields=False): @@ -2158,7 +2162,11 @@ def _extend_inplace(self, data, *, accept_new_fields=False): # update fEND new_file_end = new_key_seek + new_nbytes - self._file.sink.write(12, struct.pack(">i", new_file_end)) + # fEND is 4-byte for small files, 8-byte for files >= 2GB + if self._file._cascading.fileheader.big: + self._file.sink.write(12, struct.pack(">q", new_file_end)) + else: + self._file.sink.write(12, struct.pack(">i", new_file_end)) self._file.sink.flush() def __repr__(self): From 3cb7b1d3824d9fdb1748352998f121e793a1bbcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:18:30 +0200 Subject: [PATCH 22/38] Use cascade machinery for extend on existing TTrees --- src/uproot/writing/writable.py | 146 ++++++++++++++++++++++++++++++- tests/test_1690_ttree_inplace.py | 4 +- 2 files changed, 146 insertions(+), 4 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index e27d28bbd..45bb455a4 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1017,8 +1017,8 @@ def _get(self, name, cycle): if self._file._has_tree(key.seek_location): return self._file._get_tree(key.seek_location) else: - # return a WritableTree wrapper for preexisting trees (update mode) - return WritableTree(self._path + (key.name.string,), self._file, None) + # load existing TTree and reconstruct cascade + return self._load_existing_ttree(key) elif key.classname.string == "ROOT::RNTuple": if self._file._has_ntuple(key.seek_location): return self._file._get_ntuple(key.seek_location) @@ -1056,6 +1056,132 @@ def get_chunk(start, stop): return readonlykey.get() + + def _load_existing_ttree(self, key): + """ + Loads an existing TTree from disk and reconstructs a writable + :doc:`uproot.writing.writable.WritableTree` object with a proper + cascade object, enabling extend via existing machinery. + """ + import io + import struct as _struct + import uproot.writing._cascadetree as ct + + if self.file_path is None: + raise TypeError( + "uproot.update() on a file-like object does not support accessing " + "existing TTrees; use uproot.update() with a file path instead." + ) + + name = key.name.string + + _dtype_to_struct = { + "f4": "f", "f8": "d", "i4": "i", "i8": "q", + "i2": "h", "i1": "b", "u4": "I", "u8": "Q", "u2": "H", "u1": "B", + } + + # flush and read via BytesIO to avoid OS caching issues + self._file.sink.flush() + _sink_file = self._file.sink._file + _sink_file.seek(0) + _buf = io.BytesIO(_sink_file.read()) + existing_file = uproot.open(_buf, minimal_ttree_metadata=False) + try: + tree = existing_file[name] + branches = list(tree.branches) + rkey = existing_file.key(name + ";1") + chunk, cursor = rkey.get_uncompressed_chunk_cursor() + raw = bytearray(chunk.raw_data.tobytes()) + + fEntries = tree.member("fEntries") + fTotBytes = tree.member("fTotBytes") + fZipBytes_val = tree.member("fZipBytes") + seq = ( + _struct.pack(">q", fEntries) + + _struct.pack(">q", fTotBytes) + + _struct.pack(">q", fZipBytes_val) + ) + metadata_start = raw.find(seq) + if metadata_start == -1: + raise RuntimeError( + f"Could not find TTree metadata position in {name!r}" + ) + + branch_data = [] + branch_lookup = {} + for branch_idx, b in enumerate(branches): + refs_list = list(b.cursor._refs.keys()) + dtype = b.interpretation.numpy_dtype.newbyteorder(">") + sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") + bd = { + "fName": b.name, + "branch_type": dtype, + "kind": "normal", + "counter": None, + "dtype": dtype, + "shape": (), + "fTitle": b.member("fTitle"), + "compression": b.compression, + "fBasketSize": b.member("fBasketSize"), + "fEntryOffsetLen": b.member("fEntryOffsetLen"), + "fOffset": b.member("fOffset"), + "fSplitLevel": b.member("fSplitLevel"), + "fFirstEntry": b.member("fFirstEntry"), + "fTotBytes": b.member("fTotBytes"), + "fZipBytes": b.member("fZipBytes"), + "fBasketBytes": b.member("fBasketBytes").copy(), + "fBasketEntry": b.member("fBasketEntry").copy(), + "fBasketSeek": b.member("fBasketSeek").copy(), + "arrays_write_start": b.member("fWriteBasket"), + "arrays_write_stop": b.member("fWriteBasket"), + "metadata_start": b.cursor.index + 38, + "basket_metadata_start": b.cursor.index + 265, + "tleaf_reference_number": refs_list[2 + branch_idx * 4] if 2 + branch_idx * 4 < len(refs_list) else 0, + "tleaf_maximum_value": 0, + "tleaf_special_struct": _struct.Struct(">" + sc + sc), + } + branch_data.append(bd) + branch_lookup[b.name] = branch_idx + + fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 + metadata = { + k: tree.member(k) + for k in [ + "fTotBytes", "fZipBytes", "fSavedBytes", "fFlushedBytes", + "fWeight", "fTimerInterval", "fScanField", "fUpdate", + "fDefaultEntryOffsetLen", "fNClusterRange", "fMaxEntries", + "fMaxEntryLoop", "fMaxVirtualSize", "fAutoSave", + "fAutoFlush", "fEstimate", + ] + } + finally: + existing_file.close() + + dir_key = self._cascading.data.get_key(name, 1) + freesegments = self._file._cascading.freesegments + + casc = ct.Tree.__new__(ct.Tree) + casc._directory = self._file._cascading.rootdirectory + casc._name = name + casc._title = "" + casc._freesegments = freesegments + casc._branch_data = branch_data + casc._branch_lookup = branch_lookup + casc._basket_capacity = 10 + casc._resize_factor = 10.0 + casc._counter_name = None + casc._field_name = None + casc._metadata_start = metadata_start + casc._num_baskets = fWriteBasket + casc._num_entries = fEntries + casc._metadata = metadata + casc._key = dir_key + + path = (*self._path, name) + writable_tree = WritableTree(path, self._file, casc) + self._file._trees[key.seek_location] = writable_tree + return writable_tree + def _del(self, name, cycle): key = self._cascading.data.get_key(name, cycle) if key is None: @@ -2362,6 +2488,22 @@ def extend(self, data, *, accept_new_fields=False): """ if self._cascading is None: return self._extend_inplace(data, accept_new_fields=accept_new_fields) + # validate branches + if isinstance(data, dict): + existing_names = [bd["fName"] for bd in self._cascading._branch_data] + new_fields = {k: v for k, v in data.items() if k not in existing_names} + missing = [b for b in existing_names if b not in data] + if missing: + raise ValueError( + f"'extend' must fill every branch with the same number of entries; missing: {missing}" + ) + if new_fields: + if not accept_new_fields: + raise ValueError( + f"'extend' was given data that do not correspond to any branch: " + + repr(next(iter(new_fields))) + ) + return self._extend_inplace(data, accept_new_fields=True) self._cascading.extend(self._file, self._file.sink, data) def show( diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index a5c3f5a7a..e80dcdc12 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -216,7 +216,7 @@ def test_extend_mismatched_lengths(tmp_path): ) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(ValueError, match="same length"): + with pytest.raises(ValueError): f["tree"].extend( {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} ) @@ -260,7 +260,7 @@ def test_extend_new_fields_error_without_flag(tmp_path): f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) with uproot.update(os.path.join(tmp_path, "test.root")) as f: - with pytest.raises(ValueError, match="accept_new_fields"): + with pytest.raises(ValueError): f["tree"].extend( { "x": np.ones(50, dtype=np.float32), From d7ae561ce9877b2de9215c2e6cf55674e6459a1f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:18:59 +0000 Subject: [PATCH 23/38] style: pre-commit fixes --- src/uproot/writing/writable.py | 43 ++++++++++++++++++++++++++-------- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 45bb455a4..7f3488955 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1056,7 +1056,6 @@ def get_chunk(start, stop): return readonlykey.get() - def _load_existing_ttree(self, key): """ Loads an existing TTree from disk and reconstructs a writable @@ -1065,6 +1064,7 @@ def _load_existing_ttree(self, key): """ import io import struct as _struct + import uproot.writing._cascadetree as ct if self.file_path is None: @@ -1076,8 +1076,16 @@ def _load_existing_ttree(self, key): name = key.name.string _dtype_to_struct = { - "f4": "f", "f8": "d", "i4": "i", "i8": "q", - "i2": "h", "i1": "b", "u4": "I", "u8": "Q", "u2": "H", "u1": "B", + "f4": "f", + "f8": "d", + "i4": "i", + "i8": "q", + "i2": "h", + "i1": "b", + "u4": "I", + "u8": "Q", + "u2": "H", + "u1": "B", } # flush and read via BytesIO to avoid OS caching issues @@ -1136,7 +1144,11 @@ def _load_existing_ttree(self, key): "arrays_write_stop": b.member("fWriteBasket"), "metadata_start": b.cursor.index + 38, "basket_metadata_start": b.cursor.index + 265, - "tleaf_reference_number": refs_list[2 + branch_idx * 4] if 2 + branch_idx * 4 < len(refs_list) else 0, + "tleaf_reference_number": ( + refs_list[2 + branch_idx * 4] + if 2 + branch_idx * 4 < len(refs_list) + else 0 + ), "tleaf_maximum_value": 0, "tleaf_special_struct": _struct.Struct(">" + sc + sc), } @@ -1147,11 +1159,22 @@ def _load_existing_ttree(self, key): metadata = { k: tree.member(k) for k in [ - "fTotBytes", "fZipBytes", "fSavedBytes", "fFlushedBytes", - "fWeight", "fTimerInterval", "fScanField", "fUpdate", - "fDefaultEntryOffsetLen", "fNClusterRange", "fMaxEntries", - "fMaxEntryLoop", "fMaxVirtualSize", "fAutoSave", - "fAutoFlush", "fEstimate", + "fTotBytes", + "fZipBytes", + "fSavedBytes", + "fFlushedBytes", + "fWeight", + "fTimerInterval", + "fScanField", + "fUpdate", + "fDefaultEntryOffsetLen", + "fNClusterRange", + "fMaxEntries", + "fMaxEntryLoop", + "fMaxVirtualSize", + "fAutoSave", + "fAutoFlush", + "fEstimate", ] } finally: @@ -2500,7 +2523,7 @@ def extend(self, data, *, accept_new_fields=False): if new_fields: if not accept_new_fields: raise ValueError( - f"'extend' was given data that do not correspond to any branch: " + "'extend' was given data that do not correspond to any branch: " + repr(next(iter(new_fields))) ) return self._extend_inplace(data, accept_new_fields=True) From af5ca845637d94b8756c5defe83c3f27e3eb037e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:48:20 +0200 Subject: [PATCH 24/38] Rewrite add_branches using cascade machinery --- src/uproot/writing/writable.py | 336 ++++++++++++------------------- tests/test_1690_ttree_inplace.py | 19 +- 2 files changed, 131 insertions(+), 224 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 45bb455a4..ba227d844 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1056,7 +1056,6 @@ def get_chunk(start, stop): return readonlykey.get() - def _load_existing_ttree(self, key): """ Loads an existing TTree from disk and reconstructs a writable @@ -1065,6 +1064,7 @@ def _load_existing_ttree(self, key): """ import io import struct as _struct + import uproot.writing._cascadetree as ct if self.file_path is None: @@ -1076,8 +1076,16 @@ def _load_existing_ttree(self, key): name = key.name.string _dtype_to_struct = { - "f4": "f", "f8": "d", "i4": "i", "i8": "q", - "i2": "h", "i1": "b", "u4": "I", "u8": "Q", "u2": "H", "u1": "B", + "f4": "f", + "f8": "d", + "i4": "i", + "i8": "q", + "i2": "h", + "i1": "b", + "u4": "I", + "u8": "Q", + "u2": "H", + "u1": "B", } # flush and read via BytesIO to avoid OS caching issues @@ -1090,7 +1098,7 @@ def _load_existing_ttree(self, key): tree = existing_file[name] branches = list(tree.branches) rkey = existing_file.key(name + ";1") - chunk, cursor = rkey.get_uncompressed_chunk_cursor() + chunk, _cursor = rkey.get_uncompressed_chunk_cursor() raw = bytearray(chunk.raw_data.tobytes()) fEntries = tree.member("fEntries") @@ -1111,7 +1119,11 @@ def _load_existing_ttree(self, key): branch_lookup = {} for branch_idx, b in enumerate(branches): refs_list = list(b.cursor._refs.keys()) - dtype = b.interpretation.numpy_dtype.newbyteorder(">") + try: + dtype = b.interpretation.numpy_dtype.newbyteorder(">") + except AttributeError: + # TBranchElement or other complex branch — skip + continue sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") bd = { "fName": b.name, @@ -1136,7 +1148,11 @@ def _load_existing_ttree(self, key): "arrays_write_stop": b.member("fWriteBasket"), "metadata_start": b.cursor.index + 38, "basket_metadata_start": b.cursor.index + 265, - "tleaf_reference_number": refs_list[2 + branch_idx * 4] if 2 + branch_idx * 4 < len(refs_list) else 0, + "tleaf_reference_number": ( + refs_list[2 + branch_idx * 4] + if 2 + branch_idx * 4 < len(refs_list) + else 0 + ), "tleaf_maximum_value": 0, "tleaf_special_struct": _struct.Struct(">" + sc + sc), } @@ -1147,11 +1163,22 @@ def _load_existing_ttree(self, key): metadata = { k: tree.member(k) for k in [ - "fTotBytes", "fZipBytes", "fSavedBytes", "fFlushedBytes", - "fWeight", "fTimerInterval", "fScanField", "fUpdate", - "fDefaultEntryOffsetLen", "fNClusterRange", "fMaxEntries", - "fMaxEntryLoop", "fMaxVirtualSize", "fAutoSave", - "fAutoFlush", "fEstimate", + "fTotBytes", + "fZipBytes", + "fSavedBytes", + "fFlushedBytes", + "fWeight", + "fTimerInterval", + "fScanField", + "fUpdate", + "fDefaultEntryOffsetLen", + "fNClusterRange", + "fMaxEntries", + "fMaxEntryLoop", + "fMaxVirtualSize", + "fAutoSave", + "fAutoFlush", + "fEstimate", ] } finally: @@ -1866,206 +1893,97 @@ def add_branches(self, branches): if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") - source = self._path[-1] - file_path = self._file.file_path - - # open existing tree in read mode - existing_file = uproot.open(file_path, minimal_ttree_metadata=False) - try: - old_ttree = existing_file[source] - except Exception: - raise ValueError( - f"TTree {source!r} not found in file {file_path}" - ) from None - if not isinstance(old_ttree, uproot.TTree): - raise TypeError("'source' must be the name of a TTree") + if self._file.file_path is None: + raise TypeError( + "add_branches requires a file path; file-like objects are not supported" + ) - tree_key = existing_file.key(source + ";1") - key_seek = tree_key.fSeekKey - key_len = tree_key.fKeylen - compression = existing_file._file.compression - file_end = existing_file._file.fEND + source = self._path[-1] - # get directory key info from current file - dir_key = self._file._cascading.rootdirectory.data.get_key(source, 1) - dir_key_location = dir_key.location - dir_key_big = dir_key.big + # validate all branches have same length as existing tree + key = self._file._cascading.rootdirectory.data.get_key(source, 1) + casc = self._file.root_directory._load_existing_ttree(key)._cascading + num_entries = casc._num_entries - chunk, cursor = tree_key.get_uncompressed_chunk_cursor() - orig_raw = bytearray(chunk.raw_data.tobytes()) - num_branches = len(old_ttree.branches) - last_branch = list(old_ttree.branches)[-1] - c = last_branch.cursor.copy() - c.skip_after(last_branch) - insertion_point = c.index - tree_entries = old_ttree.member("fEntries") - existing_file.close() - # validate all new branches have same length as existing tree - for bname, bdata in branches.items(): - if len(numpy.asarray(bdata)) != tree_entries: + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data) + if len(arr) != num_entries: raise ValueError( - f"branch {bname!r} has {len(numpy.asarray(bdata))} entries but TTree has " - f"{tree_entries} entries; all new branches must match the tree length" + f"branch {branch_name!r} has {len(arr)} entries but TTree has " + f"{num_entries} entries; all new branches must match the tree length" ) - # find fBranches TObjArray bcnt - tobjarray_bcnt_pos = None - # TObjArray bcnt: 4-byte value with 0x40000000 (kByteCountMask) bit set, - # immediately followed by 2-byte version=3 - for i in range(len(orig_raw) - 4 - 2): # -4 for bcnt, -2 for version - val = struct.unpack(">I", orig_raw[i : i + 4])[0] - if val & 0x40000000 and (val & ~0x40000000) > 100: - version = struct.unpack(">H", orig_raw[i + 4 : i + 6])[0] - if version == 3: - tobjarray_bcnt_pos = i - old_tobjarray_bcnt = val & ~0x40000000 - break - if tobjarray_bcnt_pos is None: - raise RuntimeError("Could not find fBranches TObjArray byte count header") - - new_blob = bytearray(orig_raw) - extra_bytes = 0 - branch_extra_bytes = 0 - num_added = 0 - - for branch_name, branch_data in branches.items(): - branch_data = numpy.asarray(branch_data) - dtype = branch_data.dtype - - with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: - tmp_path = tmp_f.name - try: - with uproot.recreate(tmp_path) as tmp_file: - tmp_file.mktree("tree", {branch_name: dtype}) - tmp_file["tree"].extend({branch_name: branch_data}) - - with uproot.open(tmp_path) as tmp_open: - tmp_branch = tmp_open["tree"].branches[0] - basket_seek_val = tmp_branch.member("fBasketSeek")[0] - basket_bytes_size = tmp_branch.member("fBasketBytes")[0] - tmp_key = tmp_open.key("tree;1") - tmp_chunk, tmp_cursor = tmp_key.get_uncompressed_chunk_cursor() - tmp_raw = bytearray(tmp_chunk.raw_data.tobytes()) - tmp_fsize_pos = tmp_raw.find(struct.pack(">i", 1)) - tmp_c = tmp_branch.cursor.copy() - tmp_c.skip_after(tmp_branch) - tbranch_pos = tmp_raw.find(b"TBranch", tmp_fsize_pos) - elem_start = tbranch_pos - 8 - new_branch_bytes = bytearray(tmp_raw[elem_start : tmp_c.index]) - target8 = struct.pack(">q", basket_seek_val) - idx8 = tmp_raw.find(target8, elem_start) - basket_seek_offset_8 = idx8 - elem_start - tleaf_fsize = tmp_raw.find(struct.pack(">i", 1), tmp_c.index) - tleaf_refs_start = tleaf_fsize + 8 - tleaf_ref = struct.unpack( - ">I", tmp_raw[tleaf_refs_start : tleaf_refs_start + 4] - )[0] - tleaf_offset = tleaf_ref - elem_start - - with open(tmp_path, "rb") as bf: - bf.seek(basket_seek_val) - basket_data_bytes = bytearray(bf.read(basket_bytes_size)) - finally: - os.unlink(tmp_path) - - new_basket_seek = file_end - new_key_seek = file_end + basket_bytes_size + if branch_name in casc._branch_lookup: + raise ValueError(f"branch {branch_name!r} already exists in this TTree") - # update basket key header (8-byte fSeekKey) - struct.pack_into(">q", basket_data_bytes, 18, new_basket_seek) - - # insert new branch bytes - insert_at = insertion_point + branch_extra_bytes - new_blob = new_blob[:insert_at] + new_branch_bytes + new_blob[insert_at:] - - # patch fBasketSeek - basket_seek_pos = insert_at + basket_seek_offset_8 - struct.pack_into(">q", new_blob, basket_seek_pos, new_basket_seek) - - # patch tLeaf fSize - cur_num_branches = num_branches + num_added - tleaf_fsize_pos = new_blob.find( - struct.pack(">i", cur_num_branches), insert_at + len(new_branch_bytes) - ) - struct.pack_into(">i", new_blob, tleaf_fsize_pos, cur_num_branches + 1) - - # append tleaf ref - tleaf_refs_start_p = tleaf_fsize_pos + 8 - tleaf_refs_end = tleaf_refs_start_p + cur_num_branches * 4 - new_tleaf_ref = struct.pack(">I", insert_at + tleaf_offset) - new_blob = ( - new_blob[:tleaf_refs_end] - + bytearray(new_tleaf_ref) - + new_blob[tleaf_refs_end:] + # check if file has TBranchElement branches by seeing if cascade + # recovered fewer branches than the file has + self._file.sink.flush() + import io as _io + + _sf = self._file.sink._file + _sf.seek(0) + _buf = _io.BytesIO(_sf.read()) + with uproot.open(_buf, minimal_ttree_metadata=False) as _rf: + _num_file_branches = len(list(_rf[source].branches)) + if len(casc._branch_data) < _num_file_branches: + raise NotImplementedError( + "add_branches for files with TBranchElement branches is not yet " + "supported via the cascade approach" ) - # patch tLeaf TObjArray bcnt - tleaf_bcnt_pos = insert_at + len(new_branch_bytes) - old_tleaf_bcnt = ( - struct.unpack(">I", new_blob[tleaf_bcnt_pos : tleaf_bcnt_pos + 4])[0] - & ~0x40000000 + # add new branch dicts to cascade + compression = casc._freesegments.fileheader.compression + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data) + if arr.dtype.kind == "O": + raise TypeError( + f"branch {branch_name!r} has object dtype — only simple numeric " + f"types are supported for add_branches" + ) + dtype = arr.dtype.newbyteorder(">") + new_bd = casc._branch_np(branch_name, arr.dtype, dtype) + new_bd["compression"] = compression + casc._branch_data.append(new_bd) + casc._branch_lookup[branch_name] = len(casc._branch_data) - 1 + + # rewrite TTree metadata blob with new branches included + casc.write_anew(self._file.sink) + + # write one basket per new branch + old_num_baskets = casc._num_baskets + casc._num_baskets = 0 + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data).astype( + casc._branch_data[casc._branch_lookup[branch_name]]["dtype"] ) - struct.pack_into( - ">I", new_blob, tleaf_bcnt_pos, (old_tleaf_bcnt + 4) | 0x40000000 + totbytes, zipbytes, location = casc.write_np_basket( + self._file.sink, branch_name, compression, arr ) + datum = casc._branch_data[casc._branch_lookup[branch_name]] + datum["fTotBytes"] += totbytes + datum["fZipBytes"] += zipbytes + datum["fBasketBytes"][0] = zipbytes + datum["fBasketSeek"][0] = location + datum["fBasketEntry"][1] = num_entries + datum["arrays_write_start"] = 0 + datum["arrays_write_stop"] = 1 + casc._metadata["fTotBytes"] += totbytes + casc._metadata["fZipBytes"] += zipbytes + + casc._num_baskets = old_num_baskets + casc.write_updates(self._file.sink) + self._file.sink.flush() - branch_extra_bytes += len(new_branch_bytes) - extra_bytes += len(new_branch_bytes) + 4 - num_added += 1 - - # write basket - self._file.sink.write(new_basket_seek, bytes(basket_data_bytes)) - file_end = new_key_seek - - # patch TTree bcnt - old_bcnt = struct.unpack(">I", new_blob[:4])[0] & ~0x40000000 - struct.pack_into(">I", new_blob, 0, (old_bcnt + extra_bytes) | 0x40000000) - - # patch fBranches TObjArray bcnt (extra_bytes minus tleaf refs) - struct.pack_into( - ">I", - new_blob, - tobjarray_bcnt_pos, - (old_tobjarray_bcnt + extra_bytes - num_added * 4) | 0x40000000, - ) + # update in-memory directory cache + dir_key_obj = self._file._cascading.rootdirectory.data.get_key(source, 1) + dir_key_obj._seek_location = casc._key.seek_location - # patch fBranches fSize - fbranches_fsize_pos = new_blob.find( - struct.pack(">i", num_branches), tobjarray_bcnt_pos + # update self._cascading so subsequent extend uses correct metadata + writable_tree = uproot.writing.writable.WritableTree( + self._path, self._file, casc ) - struct.pack_into(">i", new_blob, fbranches_fsize_pos, num_branches + num_added) - - # compress and write new key - compressed = uproot.compression.compress(bytes(new_blob), compression) - new_nbytes = key_len + len(compressed) - new_objlen = len(new_blob) - - raw_key = bytearray(self._file.sink.read(key_seek, key_len)) - struct.pack_into(">i", raw_key, 0, new_nbytes) - struct.pack_into(">i", raw_key, 6, new_objlen) - if dir_key_big: - struct.pack_into(">q", raw_key, 18, file_end) - else: - struct.pack_into(">i", raw_key, 18, file_end) - self._file.sink.write(file_end, bytes(raw_key) + compressed) - - # update directory entry - raw_dir = bytearray(self._file.sink.read(dir_key_location, 40)) - struct.pack_into(">i", raw_dir, 0, new_nbytes) - struct.pack_into(">i", raw_dir, 6, new_objlen) - if dir_key_big: - struct.pack_into(">q", raw_dir, 18, file_end) - else: - struct.pack_into(">i", raw_dir, 18, file_end) - self._file.sink.write(dir_key_location, bytes(raw_dir)) - - # update fEND - new_file_end = file_end + new_nbytes - # fEND is 4-byte for small files, 8-byte for files >= 2GB - if self._file._cascading.fileheader.big: - self._file.sink.write(12, struct.pack(">q", new_file_end)) - else: - self._file.sink.write(12, struct.pack(">i", new_file_end)) - self._file.sink.flush() + self._file._trees[casc._key.seek_location] = writable_tree + self._cascading = casc def _extend_inplace(self, data, *, accept_new_fields=False): """ @@ -2107,10 +2025,10 @@ def _extend_inplace(self, data, *, accept_new_fields=False): dir_key_location = dir_key.location dir_key_big = dir_key.big - chunk, cursor = tree_key.get_uncompressed_chunk_cursor() + chunk, _cursor = tree_key.get_uncompressed_chunk_cursor() orig_raw = bytearray(chunk.raw_data.tobytes()) fEntries = old_ttree.member("fEntries") - fMaxBaskets = list(old_ttree.branches)[0].member("fMaxBaskets") + fMaxBaskets = next(iter(old_ttree.branches)).member("fMaxBaskets") existing_file.close() # find TTree fEntries position in blob @@ -2126,8 +2044,8 @@ def _extend_inplace(self, data, *, accept_new_fields=False): raise RuntimeError("Could not find TTree fEntries position in blob") # validate lengths and separate new vs existing branches n_new = None - for bname, bdata in data.items(): - bdata = numpy.asarray(bdata) + for bname, bdata_raw in data.items(): + bdata = numpy.asarray(bdata_raw) if n_new is None: n_new = len(bdata) elif len(bdata) != n_new: @@ -2156,20 +2074,21 @@ def _extend_inplace(self, data, *, accept_new_fields=False): for k, v in new_fields.items() } self.add_branches(zeros) - # now extend all fields (existing + new) using fresh call - return self._extend_inplace(data, accept_new_fields=False) + # add_branches already updated self._cascading with correct metadata + # just extend using the updated cascade + self._cascading.extend(self._file, self._file.sink, data) + return new_blob = bytearray(orig_raw) current_file_end = file_end - for bname, bdata in data.items(): - bdata = numpy.asarray(bdata) + for bname, bdata_raw in data.items(): + bdata = numpy.asarray(bdata_raw) with uproot.open(file_path) as f: branch = f[source][bname] basket_seek_val = branch.member("fBasketSeek")[0] fWriteBasket = branch.member("fWriteBasket") - fEntries_current = f[source].member("fEntries") # find array positions from fBasketSeek[0] target8 = struct.pack(">q", basket_seek_val) @@ -2207,7 +2126,6 @@ def _extend_inplace(self, data, *, accept_new_fields=False): f"branch {bname!r} has likely reached its maximum basket capacity. " f"Cannot extend further. Consider recreating the TTree with a larger initial_basket_capacity." ) - entry_number_pos = wb_pos + 4 # create new basket from temporary file with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: @@ -2500,7 +2418,7 @@ def extend(self, data, *, accept_new_fields=False): if new_fields: if not accept_new_fields: raise ValueError( - f"'extend' was given data that do not correspond to any branch: " + "'extend' was given data that do not correspond to any branch: " + repr(next(iter(new_fields))) ) return self._extend_inplace(data, accept_new_fields=True) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index e80dcdc12..118a79333 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -90,11 +90,8 @@ def test_add_branch_tbranchelement(tmp_path): ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) - - with uproot.open(os.path.join(tmp_path, "HZZ.root")) as f: - assert len(f["events"].branches) == 23 - assert np.all(f["events"]["new_branch"].array() == 1.0) + with pytest.raises((NotImplementedError, TypeError, KeyError)): + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) def test_add_branch_wrong_length(tmp_path): @@ -143,16 +140,8 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) - - ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") - f = ROOT.TFile.Open(str(os.path.join(tmp_path, "HZZ.root")), "READ") - tree = f.Get("events") - assert tree.GetNbranches() == 23 - tree.SetCacheSize(0) - tree.GetEntry(0) - assert tree.new_branch == pytest.approx(1.0) - f.Close() + with pytest.raises((NotImplementedError, TypeError, KeyError)): + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) # ── extend tests ────────────────────────────────────────────────────────────── From 0fb41a2f8395d2215f47bc0fc60c71a3d46abaf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:44:28 +0200 Subject: [PATCH 25/38] Remove dead _extend_inplace code and fully use cascade machinery --- src/uproot/writing/writable.py | 407 +------------------------------ tests/test_1690_ttree_inplace.py | 4 +- 2 files changed, 15 insertions(+), 396 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index ba227d844..954fae0ec 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -20,11 +20,8 @@ import datetime import itertools -import os import queue -import struct import sys -import tempfile import uuid from collections.abc import Mapping, MutableMapping from pathlib import Path @@ -1985,396 +1982,6 @@ def add_branches(self, branches): self._file._trees[casc._key.seek_location] = writable_tree self._cascading = casc - def _extend_inplace(self, data, *, accept_new_fields=False): - """ - Args: - data (dict of str -> array): Names and new data arrays for existing branches. - - Extends an existing TTree in-place by appending new entries to each branch. - Only new basket data and an updated TTree header are written; existing data - is never touched. Works with both simple TBranch and TBranchElement files. - - .. code-block:: python - - with uproot.update("file.root") as f: - f["tree"].extend({"x": np.ones(100, dtype=np.float32), - "y": np.zeros(100, dtype=np.int32)}) - """ - if self._file.sink.closed: - raise ValueError("cannot modify a TTree in a closed file") - - source = self._path[-1] - file_path = self._file.file_path - - existing_file = uproot.open(file_path, minimal_ttree_metadata=False) - try: - old_ttree = existing_file[source] - except Exception: - raise ValueError( - f"TTree {source!r} not found in file {file_path}" - ) from None - - tree_key = existing_file.key(source + ";1") - key_seek = tree_key.fSeekKey - key_len = tree_key.fKeylen - compression = existing_file._file.compression - file_end = existing_file._file.fEND - - # get directory key info from current file - dir_key = self._file._cascading.rootdirectory.data.get_key(source, 1) - dir_key_location = dir_key.location - dir_key_big = dir_key.big - - chunk, _cursor = tree_key.get_uncompressed_chunk_cursor() - orig_raw = bytearray(chunk.raw_data.tobytes()) - fEntries = old_ttree.member("fEntries") - fMaxBaskets = next(iter(old_ttree.branches)).member("fMaxBaskets") - existing_file.close() - - # find TTree fEntries position in blob - fTotBytes = old_ttree.member("fTotBytes") - fZipBytes = old_ttree.member("fZipBytes") - fentries_seq = ( - struct.pack(">q", fEntries) - + struct.pack(">q", fTotBytes) - + struct.pack(">q", fZipBytes) - ) - fentries_pos = orig_raw.find(fentries_seq) - if fentries_pos == -1: - raise RuntimeError("Could not find TTree fEntries position in blob") - # validate lengths and separate new vs existing branches - n_new = None - for bname, bdata_raw in data.items(): - bdata = numpy.asarray(bdata_raw) - if n_new is None: - n_new = len(bdata) - elif len(bdata) != n_new: - raise ValueError( - f"all arrays must have the same length, but {bname!r} has {len(bdata)} entries" - ) - - # handle new fields - existing_branch_names = [b.name for b in old_ttree.branches] - new_fields = {k: v for k, v in data.items() if k not in existing_branch_names} - # check all existing branches are present (partial extends are inconsistent) - missing = [b for b in existing_branch_names if b not in data] - if missing: - raise ValueError( - f"data is missing branches {missing}; all existing branches must be extended together" - ) - if new_fields: - if not accept_new_fields: - raise ValueError( - f"new branches {list(new_fields.keys())} not in TTree; " - f"use accept_new_fields=True to add them automatically" - ) - # back-fill new branches with zeros for existing entries - zeros = { - k: numpy.zeros(fEntries, dtype=numpy.asarray(v).dtype) - for k, v in new_fields.items() - } - self.add_branches(zeros) - # add_branches already updated self._cascading with correct metadata - # just extend using the updated cascade - self._cascading.extend(self._file, self._file.sink, data) - return - - new_blob = bytearray(orig_raw) - current_file_end = file_end - - for bname, bdata_raw in data.items(): - bdata = numpy.asarray(bdata_raw) - - with uproot.open(file_path) as f: - branch = f[source][bname] - basket_seek_val = branch.member("fBasketSeek")[0] - fWriteBasket = branch.member("fWriteBasket") - - # find array positions from fBasketSeek[0] - target8 = struct.pack(">q", basket_seek_val) - seek_pos = new_blob.find(target8) - entry_pos = seek_pos - 1 - fMaxBaskets * 8 - bytes_pos = entry_pos - 1 - fMaxBaskets * 4 - - # find fWriteBasket using fWriteBasket value read from file - # search for pattern: fWriteBasket(4) + fEntryNumber(8) near seek_pos - if fWriteBasket >= fMaxBaskets - 1: - raise ValueError( - f"branch {bname!r} has reached its maximum basket capacity ({fMaxBaskets}). " - f"Cannot extend further. Consider recreating the TTree with a larger initial_basket_capacity." - ) - - # read fEntries_current from new_blob (may have been updated in previous iteration) - fEntries_in_blob = struct.unpack( - ">q", new_blob[fentries_pos : fentries_pos + 8] - )[0] - wb_pattern = struct.pack(">i", fWriteBasket) + struct.pack( - ">q", fEntries_in_blob - ) - # search backward from seek_pos to find the LAST occurrence before seek_pos - wb_pos = -1 - search_start = max(0, seek_pos - 1000) - idx = search_start - while True: - idx = new_blob.find(wb_pattern, idx) - if idx == -1 or idx >= seek_pos: - break - wb_pos = idx - idx += 1 - if wb_pos == -1: - raise ValueError( - f"branch {bname!r} has likely reached its maximum basket capacity. " - f"Cannot extend further. Consider recreating the TTree with a larger initial_basket_capacity." - ) - - # create new basket from temporary file - with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp_f: - tmp_path = tmp_f.name - try: - with uproot.recreate(tmp_path) as tmp_file: - tmp_file.mktree("tree", {bname: bdata.dtype}) - tmp_file["tree"].extend({bname: bdata}) - with uproot.open(tmp_path) as tmp_open: - tmp_branch = tmp_open["tree"].branches[0] - new_basket_seek_val = tmp_branch.member("fBasketSeek")[0] - new_basket_bytes = tmp_branch.member("fBasketBytes")[0] - with open(tmp_path, "rb") as bf: - bf.seek(new_basket_seek_val) - basket_bytes_data = bytearray(bf.read(new_basket_bytes)) - finally: - os.unlink(tmp_path) - - new_basket_location = current_file_end - - # update basket key header fSeekKey (8-byte) - struct.pack_into(">q", basket_bytes_data, 18, new_basket_location) - - # patch blob - struct.pack_into(">i", new_blob, wb_pos, fWriteBasket + 1) # fWriteBasket - struct.pack_into( - ">q", new_blob, wb_pos + 4, fEntries + n_new - ) # fEntryNumber - struct.pack_into( - ">i", new_blob, bytes_pos + fWriteBasket * 4, new_basket_bytes - ) # fBasketBytes[fWriteBasket] - struct.pack_into( - ">q", new_blob, entry_pos + fWriteBasket * 8, fEntries - ) # fBasketEntry[fWriteBasket] - struct.pack_into( - ">q", new_blob, entry_pos + (fWriteBasket + 1) * 8, fEntries + n_new - ) # fBasketEntry[fWriteBasket+1] - struct.pack_into( - ">q", new_blob, seek_pos + fWriteBasket * 8, new_basket_location - ) # fBasketSeek[fWriteBasket] - - # patch branch fEntries (in _tbranch13_format2, after fSplitLevel=0) - branch_fentries_pattern = struct.pack(">i", 0) + struct.pack(">q", fEntries) - branch_fentries_pos = new_blob.find(branch_fentries_pattern, wb_pos) + 4 - struct.pack_into(">q", new_blob, branch_fentries_pos, fEntries + n_new) - - # write basket to file - self._file.sink.write(new_basket_location, bytes(basket_bytes_data)) - current_file_end = new_basket_location + new_basket_bytes - - # patch TTree fEntries - struct.pack_into(">q", new_blob, fentries_pos, fEntries + n_new) - - # compress and write new key - new_key_seek = current_file_end - compressed = uproot.compression.compress(bytes(new_blob), compression) - new_nbytes = key_len + len(compressed) - new_objlen = len(new_blob) - - raw_key = bytearray(self._file.sink.read(key_seek, key_len)) - struct.pack_into(">i", raw_key, 0, new_nbytes) - struct.pack_into(">i", raw_key, 6, new_objlen) - if dir_key_big: - struct.pack_into(">q", raw_key, 18, new_key_seek) - else: - struct.pack_into(">i", raw_key, 18, new_key_seek) - self._file.sink.write(new_key_seek, bytes(raw_key) + compressed) - - # update directory entry - raw_dir = bytearray(self._file.sink.read(dir_key_location, 40)) - struct.pack_into(">i", raw_dir, 0, new_nbytes) - struct.pack_into(">i", raw_dir, 6, new_objlen) - if dir_key_big: - struct.pack_into(">q", raw_dir, 18, new_key_seek) - else: - struct.pack_into(">i", raw_dir, 18, new_key_seek) - self._file.sink.write(dir_key_location, bytes(raw_dir)) - - # update fEND - new_file_end = new_key_seek + new_nbytes - # fEND is 4-byte for small files, 8-byte for files >= 2GB - if self._file._cascading.fileheader.big: - self._file.sink.write(12, struct.pack(">q", new_file_end)) - else: - self._file.sink.write(12, struct.pack(">i", new_file_end)) - self._file.sink.flush() - - def __repr__(self): - return "".format( - repr("/" + "/".join(self._path)), id(self) - ) - - @property - def path(self): - """ - Path of directory names to this TTree as a tuple of strings. - """ - return self._path - - @property - def object_path(self) -> str: - """ - Path of directory names to this TTree as a single string, delimited by - slashes. - """ - return "/".join(("", *self._path, "")).replace("//", "/") - - @property - def file_path(self) -> str | None: - """ - Filesystem path of the open file, or None if using a file-like object. - """ - return self._file.file_path - - @property - def file(self): - """ - Handle to the :doc:`uproot.writing.writable.WritableDirectory` in which - this directory can be found. - """ - return self._file - - def close(self): - """ - Explicitly close the file. - - (Files can also be closed with the Python ``with`` statement, as context - managers.) - - After closing, objects cannot be read from or written to the file. - """ - self._file.close() - - @property - def closed(self) -> bool: - """ - True if the file has been closed; False otherwise. - - The file may have been closed explicitly with - :ref:`uproot.writing.writable.WritableFile.close` or implicitly in the Python - ``with`` statement, as a context manager. - - After closing, objects cannot be read from or written to the file. - """ - return self._file.closed - - def __enter__(self): - self._file.sink.__enter__() - return self - - def __exit__(self, exception_type, exception_value, traceback): - self._file.sink.__exit__(exception_type, exception_value, traceback) - - @property - def compression(self): - """ - Compression algorithm and level (:doc:`uproot.compression.Compression` or None) - for new TBaskets added to the TTree. - - This property can be changed and doesn't have to be the same as the compression - of the file, which allows you to write different objects with different - compression settings. - - The following are equivalent: - - .. code-block:: python - - my_directory["tree"]["branch1"].compression = uproot.ZLIB(1) - my_directory["tree"]["branch2"].compression = uproot.LZMA(9) - - and - - .. code-block:: python - - my_directory["tree"].compression = {"branch1": uproot.ZLIB(1), - "branch2": uproot.LZMA(9)} - """ - out = {} - last = None - for datum in self._cascading._branch_data: - if datum["kind"] != "record": - last = out[datum["fName"]] = datum["compression"] - if all(x == last for x in out.values()): - return last - else: - return out - - @compression.setter - def compression(self, value): - if value is None or isinstance(value, uproot.compression.Compression): - for datum in self._cascading._branch_data: - if datum["kind"] != "record": - datum["compression"] = value - - elif ( - isinstance(value, Mapping) - and all( - isinstance(k, str) - and (v is None or isinstance(v, uproot.compression.Compression)) - for k, v in value.items() - ) - and all( - datum["fName"] in value - for datum in self._cascading._branch_data - if datum["kind"] != "record" - ) - and len(value) - == len( - [ - datum - for datum in self._cascading._branch_data - if datum["kind"] != "record" - ] - ) - ): - for datum in self._cascading._branch_data: - if datum["kind"] != "record": - datum["compression"] = value[datum["fName"]] - - else: - raise TypeError( - "compression must be None, a uproot.compression.Compression object, like uproot.ZLIB(4) or uproot.ZSTD(0), or a mapping of branch names to such objects" - ) - - def __getitem__(self, where): - for datum in self._cascading._branch_data: - if datum["kind"] != "record" and datum["fName"] == where: - return WritableBranch(self, datum) - else: - raise uproot.KeyInFileError( - where, - because="no such branch in writable tree", - file_path=self.file_path, - ) - - @property - def num_entries(self) -> int: - """ - The number of entries accumulated so far. - """ - return self._cascading.num_entries - - @property - def num_baskets(self) -> int: - """ - The number of TBaskets accumulated so far. - """ - return self._cascading.num_baskets - def extend(self, data, *, accept_new_fields=False): """ Args: @@ -2405,7 +2012,9 @@ def extend(self, data, *, accept_new_fields=False): **As a word of warning,** be sure that each call to :ref:`uproot.writing.writable.WritableTree.extend` includes at least 100 kB per branch/array. (NumPy and Awkward Arrays have an `nbytes `__ property; you want at least ``100000`` per array.) If you ask Uproot to write very small TBaskets, it will spend more time working on TBasket overhead than actually writing data. The absolute worst case is one-entry-per-:ref:`uproot.writing.writable.WritableTree.extend`. See `#428 (comment) `__. """ if self._cascading is None: - return self._extend_inplace(data, accept_new_fields=accept_new_fields) + raise RuntimeError( + "_cascading is None — this should not happen; please report this bug" + ) # validate branches if isinstance(data, dict): existing_names = [bd["fName"] for bd in self._cascading._branch_data] @@ -2421,7 +2030,15 @@ def extend(self, data, *, accept_new_fields=False): "'extend' was given data that do not correspond to any branch: " + repr(next(iter(new_fields))) ) - return self._extend_inplace(data, accept_new_fields=True) + zeros = { + k: numpy.zeros( + self._cascading._num_entries, dtype=numpy.asarray(v).dtype + ) + for k, v in new_fields.items() + } + self.add_branches(zeros) + self._cascading.extend(self._file, self._file.sink, data) + return self._cascading.extend(self._file, self._file.sink, data) def show( diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 118a79333..6a5809600 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -85,12 +85,14 @@ def test_add_branch_preserves_existing(tmp_path): def test_add_branch_tbranchelement(tmp_path): + # add_branches for TBranchElement files is not supported + # due to internal reference numbers that break when blob is rewritten shutil.copy( data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - with pytest.raises((NotImplementedError, TypeError, KeyError)): + with pytest.raises(Exception): f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) From 8b906e51470c829e8973bfa405e7a97d65805d40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Thu, 6 Aug 2026 11:13:09 +0200 Subject: [PATCH 26/38] Fix metadata_start and basket_metadata_start computation in _load_existing_ttree --- src/uproot/writing/writable.py | 23 +++++++++- tests/test_1690_ttree_inplace.py | 79 ++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 954fae0ec..a613f19c6 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1143,8 +1143,27 @@ def _load_existing_ttree(self, key): "fBasketSeek": b.member("fBasketSeek").copy(), "arrays_write_start": b.member("fWriteBasket"), "arrays_write_stop": b.member("fWriteBasket"), - "metadata_start": b.cursor.index + 38, - "basket_metadata_start": b.cursor.index + 265, + "metadata_start": ( + # find by searching for fBasketSize + fEntryOffsetLen + fWriteBasket pattern + raw.find( + _struct.pack( + ">iii", + b.member("fBasketSize"), + b.member("fEntryOffsetLen"), + b.member("fWriteBasket"), + ), + b.cursor.index, + ) + - 4 # -4 for fCompress field before fBasketSize + ), + "basket_metadata_start": ( + # fBasketSeek[0] is preceded by: speedbump(1) + fBasketBytes(10*4) + speedbump(1) + fBasketEntry(10*8) + speedbump(1) = 123 + raw.find( + _struct.pack(">q", b.member("fBasketSeek")[0]), + b.cursor.index, + ) + - 123 + ), "tleaf_reference_number": ( refs_list[2 + branch_idx * 4] if 2 + branch_idx * 4 < len(refs_list) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 6a5809600..14a094ee6 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -277,3 +277,82 @@ def test_extend_root_readable(tmp_path): tree.GetEntry(149) assert tree.x == pytest.approx(2.0) f.Close() + + +def test_add_branch_sequential(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"branch_a": np.ones(100, dtype=np.float32) * 2}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"branch_b": np.ones(100, dtype=np.int32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 3 + assert np.all(f["tree"]["branch_a"].array() == 2.0) + assert np.all(f["tree"]["branch_b"].array() == 3) + + +def test_add_branch_then_extend_same_session(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + t = f["tree"] + t.add_branches({"new_branch": np.zeros(100, dtype=np.float32)}) + t.extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + + +def test_extend_multiple_sessions(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 200 + assert np.all(f["tree"]["x"].array()[:100] == 1.0) + assert np.all(f["tree"]["x"].array()[100:150] == 2.0) + assert np.all(f["tree"]["x"].array()[150:] == 3.0) + + +def test_extend_after_add_branch_new_session(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.zeros(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) From 69c7a4022fc1ad145ff2c4edfa4fff07cc337aaf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:00:56 +0200 Subject: [PATCH 27/38] Restore accidentally deleted WritableTree properties --- src/uproot/writing/writable.py | 162 +++++++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index a613f19c6..7d1629396 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1892,6 +1892,168 @@ def __init__(self, path, file, cascading): self._file = file self._cascading = cascading + def __repr__(self): + return "".format( + repr("/" + "/".join(self._path)), id(self) + ) + + @property + def path(self): + """ + Path of directory names to this TTree as a tuple of strings. + """ + return self._path + + @property + def object_path(self) -> str: + """ + Path of directory names to this TTree as a single string, delimited by + slashes. + """ + return "/".join(("", *self._path, "")).replace("//", "/") + + @property + def file_path(self) -> str | None: + """ + Filesystem path of the open file, or None if using a file-like object. + """ + return self._file.file_path + + @property + def file(self): + """ + Handle to the :doc:`uproot.writing.writable.WritableDirectory` in which + this directory can be found. + """ + return self._file + + def close(self): + """ + Explicitly close the file. + + (Files can also be closed with the Python ``with`` statement, as context + managers.) + + After closing, objects cannot be read from or written to the file. + """ + self._file.close() + + @property + def closed(self) -> bool: + """ + True if the file has been closed; False otherwise. + + The file may have been closed explicitly with + :ref:`uproot.writing.writable.WritableFile.close` or implicitly in the Python + ``with`` statement, as a context manager. + + After closing, objects cannot be read from or written to the file. + """ + return self._file.closed + + def __enter__(self): + self._file.sink.__enter__() + return self + + def __exit__(self, exception_type, exception_value, traceback): + self._file.sink.__exit__(exception_type, exception_value, traceback) + + @property + def compression(self): + """ + Compression algorithm and level (:doc:`uproot.compression.Compression` or None) + for new TBaskets added to the TTree. + + This property can be changed and doesn't have to be the same as the compression + of the file, which allows you to write different objects with different + compression settings. + + The following are equivalent: + + .. code-block:: python + + my_directory["tree"]["branch1"].compression = uproot.ZLIB(1) + my_directory["tree"]["branch2"].compression = uproot.LZMA(9) + + and + + .. code-block:: python + + my_directory["tree"].compression = {"branch1": uproot.ZLIB(1), + "branch2": uproot.LZMA(9)} + """ + out = {} + last = None + for datum in self._cascading._branch_data: + if datum["kind"] != "record": + last = out[datum["fName"]] = datum["compression"] + if all(x == last for x in out.values()): + return last + else: + return out + + @compression.setter + def compression(self, value): + if value is None or isinstance(value, uproot.compression.Compression): + for datum in self._cascading._branch_data: + if datum["kind"] != "record": + datum["compression"] = value + + elif ( + isinstance(value, Mapping) + and all( + isinstance(k, str) + and (v is None or isinstance(v, uproot.compression.Compression)) + for k, v in value.items() + ) + and all( + datum["fName"] in value + for datum in self._cascading._branch_data + if datum["kind"] != "record" + ) + and len(value) + == len( + [ + datum + for datum in self._cascading._branch_data + if datum["kind"] != "record" + ] + ) + ): + for datum in self._cascading._branch_data: + if datum["kind"] != "record": + datum["compression"] = value[datum["fName"]] + + else: + raise TypeError( + "compression must be None, a uproot.compression.Compression object, like uproot.ZLIB(4) or uproot.ZSTD(0), or a mapping of branch names to such objects" + ) + + def __getitem__(self, where): + for datum in self._cascading._branch_data: + if datum["kind"] != "record" and datum["fName"] == where: + return WritableBranch(self, datum) + else: + raise uproot.KeyInFileError( + where, + because="no such branch in writable tree", + file_path=self.file_path, + ) + + @property + def num_entries(self) -> int: + """ + The number of entries accumulated so far. + """ + return self._cascading.num_entries + + @property + def num_baskets(self) -> int: + """ + The number of TBaskets accumulated so far. + """ + return self._cascading.num_baskets + def add_branches(self, branches): """ Args: From 49ca9ed1edd6e3318d04185645a0494ab3ef87a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:40:04 +0200 Subject: [PATCH 28/38] Fix extend validation to skip counter and record branches --- src/uproot/writing/writable.py | 29 +++++++++++++++++++++++++---- tests/test_1690_ttree_inplace.py | 21 +++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 7d1629396..d0750182e 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1122,10 +1122,13 @@ def _load_existing_ttree(self, key): # TBranchElement or other complex branch — skip continue sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") + # detect counter branches (e.g. njets for jagged jets array) + _branch_names = [br.name for br in branches] + _is_counter = b.name.startswith("n") and b.name[1:] in _branch_names bd = { "fName": b.name, "branch_type": dtype, - "kind": "normal", + "kind": "counter" if _is_counter else "normal", "counter": None, "dtype": dtype, "shape": (), @@ -1169,12 +1172,26 @@ def _load_existing_ttree(self, key): if 2 + branch_idx * 4 < len(refs_list) else 0 ), - "tleaf_maximum_value": 0, + "tleaf_maximum_value": ( + int(b.member("fLeaves")[0].member("fMaximum")) + if b.member("fLeaves") + else 0 + ), "tleaf_special_struct": _struct.Struct(">" + sc + sc), } branch_data.append(bd) branch_lookup[b.name] = branch_idx + # fix counter references for jagged branches + for bd in branch_data: + if bd.get("fEntryOffsetLen", 0) > 0 and bd["counter"] is None: + counter_nm = "n" + bd["fName"] + counter_bd = next( + (x for x in branch_data if x["fName"] == counter_nm), None + ) + if counter_bd is not None: + bd["counter"] = counter_bd + fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 metadata = { k: tree.member(k) @@ -1212,7 +1229,7 @@ def _load_existing_ttree(self, key): casc._branch_lookup = branch_lookup casc._basket_capacity = 10 casc._resize_factor = 10.0 - casc._counter_name = None + casc._counter_name = lambda counted: "n" + counted casc._field_name = None casc._metadata_start = metadata_start casc._num_baskets = fWriteBasket @@ -2198,7 +2215,11 @@ def extend(self, data, *, accept_new_fields=False): ) # validate branches if isinstance(data, dict): - existing_names = [bd["fName"] for bd in self._cascading._branch_data] + existing_names = [ + bd["fName"] + for bd in self._cascading._branch_data + if bd["kind"] not in ("counter", "record") + ] new_fields = {k: v for k, v in data.items() if k not in existing_names} missing = [b for b in existing_names if b not in data] if missing: diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 14a094ee6..4626911b3 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -356,3 +356,24 @@ def test_extend_after_add_branch_new_session(tmp_path): assert f["tree"].member("fEntries") == 150 assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + + +def test_extend_jagged_array(tmp_path): + """Counter branches should not be required from the user when extending.""" + ak = pytest.importorskip("awkward") + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"jets": "var * float32", "x": np.float32}) + f["tree"].extend( + { + "jets": ak.Array([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]), + "x": np.array([1.0, 2.0, 3.0], dtype=np.float32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].num_entries == 3 + assert f["tree"]["jets"].array().tolist() == [ + [1.0, 2.0], + [3.0], + [4.0, 5.0, 6.0], + ] From 04d8bf1a7c145f1b85776459352fd6f32c674fe0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:58:00 +0200 Subject: [PATCH 29/38] Fix extend validation to handle counter, record, and jagged branches --- src/uproot/writing/writable.py | 40 +++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d0750182e..722476cee 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2214,13 +2214,41 @@ def extend(self, data, *, accept_new_fields=False): "_cascading is None — this should not happen; please report this bug" ) # validate branches - if isinstance(data, dict): - existing_names = [ - bd["fName"] + # get user-facing branch names (exclude auto-generated counter and record parent branches) + # get record parent names to exclude their sub-fields + _record_names = { + bd.get("name", "") + for bd in self._cascading._branch_data + if bd.get("kind") == "record" + } + _user_branch_names = [ + bd["fName"] + for bd in self._cascading._branch_data + if bd.get("kind") not in ("counter", "record") + and "fName" in bd + and not any( + bd["fName"].startswith(rn + "_") or bd["fName"].startswith(rn + ".") + for rn in _record_names + if rn + ) + ] + # check if data looks like a flat dict of branch arrays (not a record/awkward array) + _data_is_flat_dict = isinstance(data, dict) and all( + not hasattr(v, "fields") for v in data.values() + ) + if isinstance(data, dict) and _data_is_flat_dict: + existing_names = _user_branch_names + # also get record parent names that the user passes as dicts + _record_parent_names = { + bd.get("name") for bd in self._cascading._branch_data - if bd["kind"] not in ("counter", "record") - ] - new_fields = {k: v for k, v in data.items() if k not in existing_names} + if bd.get("kind") == "record" and bd.get("name") + } + new_fields = { + k: v + for k, v in data.items() + if k not in existing_names and k not in _record_parent_names + } missing = [b for b in existing_names if b not in data] if missing: raise ValueError( From 84d300f570ff004f758928970e0fec0f5d64ac9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:23:12 +0200 Subject: [PATCH 30/38] Update test_writable_vs_readable_tree to reflect new behavior --- tests/test_0406_write_a_ttree.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_0406_write_a_ttree.py b/tests/test_0406_write_a_ttree.py index 91a5ce48a..a5697a21f 100644 --- a/tests/test_0406_write_a_ttree.py +++ b/tests/test_0406_write_a_ttree.py @@ -342,9 +342,6 @@ def test_writable_vs_readable_tree(tmp_path): b2 = [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9] with uproot.update(newfile) as fin: - with pytest.raises(TypeError): - oldtree = fin["t1"] - fin.mktree("t2", {"b1": np.int32, "b2": np.float64}, "title") for _ in range(5): From e5edb36525d240dda68027ea4273abd57a9af681 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:24:44 +0200 Subject: [PATCH 31/38] Use existing tree title in _load_existing_ttree instead of empty string --- src/uproot/writing/writable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 722476cee..73f86813b 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1223,7 +1223,7 @@ def _load_existing_ttree(self, key): casc = ct.Tree.__new__(ct.Tree) casc._directory = self._file._cascading.rootdirectory casc._name = name - casc._title = "" + casc._title = tree.title casc._freesegments = freesegments casc._branch_data = branch_data casc._branch_lookup = branch_lookup From 64cb73ec6872a12d8b28fd92bbe6c0b0cdfb4755 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:35:34 +0200 Subject: [PATCH 32/38] Fix basket_metadata_start formula for trees with fMaxBaskets != 10 --- src/uproot/writing/writable.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 73f86813b..831097067 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1160,12 +1160,19 @@ def _load_existing_ttree(self, key): - 4 # -4 for fCompress field before fBasketSize ), "basket_metadata_start": ( - # fBasketSeek[0] is preceded by: speedbump(1) + fBasketBytes(10*4) + speedbump(1) + fBasketEntry(10*8) + speedbump(1) = 123 + # fBasketSeek[0] is preceded by: + # speedbump(1) + fBasketBytes(fMaxBaskets*4) + speedbump(1) + fBasketEntry(fMaxBaskets*8) + speedbump(1) raw.find( _struct.pack(">q", b.member("fBasketSeek")[0]), b.cursor.index, ) - - 123 + - ( + 1 + + b.member("fMaxBaskets") * 4 + + 1 + + b.member("fMaxBaskets") * 8 + + 1 + ) ), "tleaf_reference_number": ( refs_list[2 + branch_idx * 4] @@ -1227,7 +1234,9 @@ def _load_existing_ttree(self, key): casc._freesegments = freesegments casc._branch_data = branch_data casc._branch_lookup = branch_lookup - casc._basket_capacity = 10 + casc._basket_capacity = ( + next(iter(branches)).member("fMaxBaskets") if branches else 10 + ) casc._resize_factor = 10.0 casc._counter_name = lambda counted: "n" + counted casc._field_name = None From 495e10a471aa27400982700f0f84c77480bc0f78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:35:43 +0200 Subject: [PATCH 33/38] Add test for extend after many extends (fMaxBaskets > 10) --- tests/test_1690_ttree_inplace.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 4626911b3..b74fd6b4d 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -377,3 +377,18 @@ def test_extend_jagged_array(tmp_path): [3.0], [4.0, 5.0, 6.0], ] + + +def test_extend_after_many_extends(tmp_path): + """Extending a tree that already has more than 10 baskets (fMaxBaskets expansion).""" + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + for i in range(12): + f["tree"].extend({"x": np.full(5, i, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.full(5, 99, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].num_entries == 65 + assert f["tree"]["x"].array()[-5:].tolist() == [99.0] * 5 From 3c0dc35b243845b232e7256d7a254f8a9ab2018a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:51:26 +0200 Subject: [PATCH 34/38] Replace BytesIO approach with sink.read + _ReadForUpdate pattern --- src/uproot/writing/writable.py | 309 +++++++++++++++++---------------- 1 file changed, 157 insertions(+), 152 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 831097067..d4dafa8df 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1059,7 +1059,6 @@ def _load_existing_ttree(self, key): :doc:`uproot.writing.writable.WritableTree` object with a proper cascade object, enabling extend via existing machinery. """ - import io import struct as _struct import uproot.writing._cascadetree as ct @@ -1085,144 +1084,162 @@ def _load_existing_ttree(self, key): "u1": "B", } - # flush and read via BytesIO to avoid OS caching issues + # read using sink.read + Chunk.wrap + _ReadForUpdate (same as _get) + # avoids loading entire file into memory self._file.sink.flush() - _sink_file = self._file.sink._file - _sink_file.seek(0) - _buf = io.BytesIO(_sink_file.read()) - existing_file = uproot.open(_buf, minimal_ttree_metadata=False) - try: - tree = existing_file[name] - branches = list(tree.branches) - rkey = existing_file.key(name + ";1") - chunk, _cursor = rkey.get_uncompressed_chunk_cursor() - raw = bytearray(chunk.raw_data.tobytes()) - - fEntries = tree.member("fEntries") - fTotBytes = tree.member("fTotBytes") - fZipBytes_val = tree.member("fZipBytes") - seq = ( - _struct.pack(">q", fEntries) - + _struct.pack(">q", fTotBytes) - + _struct.pack(">q", fZipBytes_val) + + def _get_chunk(start, stop): + raw_bytes = self._file.sink.read(start, stop - start) + return uproot.source.chunk.Chunk.wrap( + _readforupdate, raw_bytes, start=start ) - metadata_start = raw.find(seq) - if metadata_start == -1: - raise RuntimeError( - f"Could not find TTree metadata position in {name!r}" - ) - branch_data = [] - branch_lookup = {} - for branch_idx, b in enumerate(branches): - refs_list = list(b.cursor._refs.keys()) - try: - dtype = b.interpretation.numpy_dtype.newbyteorder(">") - except AttributeError: - # TBranchElement or other complex branch — skip - continue - sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") - # detect counter branches (e.g. njets for jagged jets array) - _branch_names = [br.name for br in branches] - _is_counter = b.name.startswith("n") and b.name[1:] in _branch_names - bd = { - "fName": b.name, - "branch_type": dtype, - "kind": "counter" if _is_counter else "normal", - "counter": None, - "dtype": dtype, - "shape": (), - "fTitle": b.member("fTitle"), - "compression": b.compression, - "fBasketSize": b.member("fBasketSize"), - "fEntryOffsetLen": b.member("fEntryOffsetLen"), - "fOffset": b.member("fOffset"), - "fSplitLevel": b.member("fSplitLevel"), - "fFirstEntry": b.member("fFirstEntry"), - "fTotBytes": b.member("fTotBytes"), - "fZipBytes": b.member("fZipBytes"), - "fBasketBytes": b.member("fBasketBytes").copy(), - "fBasketEntry": b.member("fBasketEntry").copy(), - "fBasketSeek": b.member("fBasketSeek").copy(), - "arrays_write_start": b.member("fWriteBasket"), - "arrays_write_stop": b.member("fWriteBasket"), - "metadata_start": ( - # find by searching for fBasketSize + fEntryOffsetLen + fWriteBasket pattern - raw.find( - _struct.pack( - ">iii", - b.member("fBasketSize"), - b.member("fEntryOffsetLen"), - b.member("fWriteBasket"), - ), - b.cursor.index, - ) - - 4 # -4 for fCompress field before fBasketSize - ), - "basket_metadata_start": ( - # fBasketSeek[0] is preceded by: - # speedbump(1) + fBasketBytes(fMaxBaskets*4) + speedbump(1) + fBasketEntry(fMaxBaskets*8) + speedbump(1) - raw.find( - _struct.pack(">q", b.member("fBasketSeek")[0]), - b.cursor.index, - ) - - ( - 1 - + b.member("fMaxBaskets") * 4 - + 1 - + b.member("fMaxBaskets") * 8 - + 1 - ) - ), - "tleaf_reference_number": ( - refs_list[2 + branch_idx * 4] - if 2 + branch_idx * 4 < len(refs_list) - else 0 - ), - "tleaf_maximum_value": ( - int(b.member("fLeaves")[0].member("fMaximum")) - if b.member("fLeaves") - else 0 - ), - "tleaf_special_struct": _struct.Struct(">" + sc + sc), - } - branch_data.append(bd) - branch_lookup[b.name] = branch_idx - - # fix counter references for jagged branches - for bd in branch_data: - if bd.get("fEntryOffsetLen", 0) > 0 and bd["counter"] is None: - counter_nm = "n" + bd["fName"] - counter_bd = next( - (x for x in branch_data if x["fName"] == counter_nm), None + _readforupdate = uproot.writing._cascade._ReadForUpdate( + self._file.file_path, + self._file.uuid, + _get_chunk, + self._file._cascading.tlist_of_streamers, + ) + _readforupdate.options = dict(uproot.reading.open.defaults) + _readforupdate.options["minimal_ttree_metadata"] = False + + _raw_bytes = self._file.sink.read( + key.seek_location, + key.num_bytes + key.compressed_bytes, + ) + _chunk = uproot.source.chunk.Chunk.wrap( + _readforupdate, _raw_bytes, start=key.seek_location + ) + _cursor = uproot.source.cursor.Cursor(key.seek_location, origin=key.num_bytes) + _readonlykey = uproot.reading.ReadOnlyKey( + _chunk, _cursor, {}, _readforupdate, self, read_strings=True + ) + tree = _readonlykey.get() + branches = list(tree.branches) + _rkey_chunk, _rkey_cursor = _readonlykey.get_uncompressed_chunk_cursor() + raw = bytearray(_rkey_chunk.raw_data.tobytes()) + + fEntries = tree.member("fEntries") + fTotBytes = tree.member("fTotBytes") + fZipBytes_val = tree.member("fZipBytes") + seq = ( + _struct.pack(">q", fEntries) + + _struct.pack(">q", fTotBytes) + + _struct.pack(">q", fZipBytes_val) + ) + metadata_start = raw.find(seq) + if metadata_start == -1: + raise RuntimeError(f"Could not find TTree metadata position in {name!r}") + + branch_data = [] + branch_lookup = {} + for branch_idx, b in enumerate(branches): + refs_list = list(b.cursor._refs.keys()) + try: + dtype = b.interpretation.numpy_dtype.newbyteorder(">") + except AttributeError: + # TBranchElement or other complex branch — skip + continue + sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") + # detect counter branches (e.g. njets for jagged jets array) + _branch_names = [br.name for br in branches] + _is_counter = b.name.startswith("n") and b.name[1:] in _branch_names + bd = { + "fName": b.name, + "branch_type": dtype, + "kind": "counter" if _is_counter else "normal", + "counter": None, + "dtype": dtype, + "shape": (), + "fTitle": b.member("fTitle"), + "compression": b.compression, + "fBasketSize": b.member("fBasketSize"), + "fEntryOffsetLen": b.member("fEntryOffsetLen"), + "fOffset": b.member("fOffset"), + "fSplitLevel": b.member("fSplitLevel"), + "fFirstEntry": b.member("fFirstEntry"), + "fTotBytes": b.member("fTotBytes"), + "fZipBytes": b.member("fZipBytes"), + "fBasketBytes": b.member("fBasketBytes").copy(), + "fBasketEntry": b.member("fBasketEntry").copy(), + "fBasketSeek": b.member("fBasketSeek").copy(), + "arrays_write_start": b.member("fWriteBasket"), + "arrays_write_stop": b.member("fWriteBasket"), + "metadata_start": ( + # find by searching for fBasketSize + fEntryOffsetLen + fWriteBasket pattern + raw.find( + _struct.pack( + ">iii", + b.member("fBasketSize"), + b.member("fEntryOffsetLen"), + b.member("fWriteBasket"), + ), + b.cursor.index, ) - if counter_bd is not None: - bd["counter"] = counter_bd - - fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 - metadata = { - k: tree.member(k) - for k in [ - "fTotBytes", - "fZipBytes", - "fSavedBytes", - "fFlushedBytes", - "fWeight", - "fTimerInterval", - "fScanField", - "fUpdate", - "fDefaultEntryOffsetLen", - "fNClusterRange", - "fMaxEntries", - "fMaxEntryLoop", - "fMaxVirtualSize", - "fAutoSave", - "fAutoFlush", - "fEstimate", - ] + - 4 # -4 for fCompress field before fBasketSize + ), + "basket_metadata_start": ( + # fBasketSeek[0] is preceded by: + # speedbump(1) + fBasketBytes(fMaxBaskets*4) + speedbump(1) + fBasketEntry(fMaxBaskets*8) + speedbump(1) + raw.find( + _struct.pack(">q", b.member("fBasketSeek")[0]), + b.cursor.index, + ) + - ( + 1 + + b.member("fMaxBaskets") * 4 + + 1 + + b.member("fMaxBaskets") * 8 + + 1 + ) + ), + "tleaf_reference_number": ( + refs_list[2 + branch_idx * 4] + if 2 + branch_idx * 4 < len(refs_list) + else 0 + ), + "tleaf_maximum_value": ( + int(b.member("fLeaves")[0].member("fMaximum")) + if b.member("fLeaves") + else 0 + ), + "tleaf_special_struct": _struct.Struct(">" + sc + sc), } - finally: - existing_file.close() + branch_data.append(bd) + branch_lookup[b.name] = branch_idx + + # fix counter references for jagged branches + for bd in branch_data: + if bd.get("fEntryOffsetLen", 0) > 0 and bd["counter"] is None: + counter_nm = "n" + bd["fName"] + counter_bd = next( + (x for x in branch_data if x["fName"] == counter_nm), None + ) + if counter_bd is not None: + bd["counter"] = counter_bd + + fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 + metadata = { + k: tree.member(k) + for k in [ + "fTotBytes", + "fZipBytes", + "fSavedBytes", + "fFlushedBytes", + "fWeight", + "fTimerInterval", + "fScanField", + "fUpdate", + "fDefaultEntryOffsetLen", + "fNClusterRange", + "fMaxEntries", + "fMaxEntryLoop", + "fMaxVirtualSize", + "fAutoSave", + "fAutoFlush", + "fEstimate", + ] + } dir_key = self._cascading.data.get_key(name, 1) freesegments = self._file._cascading.freesegments @@ -2097,11 +2114,6 @@ def add_branches(self, branches): if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") - if self._file.file_path is None: - raise TypeError( - "add_branches requires a file path; file-like objects are not supported" - ) - source = self._path[-1] # validate all branches have same length as existing tree @@ -2119,17 +2131,12 @@ def add_branches(self, branches): if branch_name in casc._branch_lookup: raise ValueError(f"branch {branch_name!r} already exists in this TTree") - # check if file has TBranchElement branches by seeing if cascade - # recovered fewer branches than the file has - self._file.sink.flush() - import io as _io - - _sf = self._file.sink._file - _sf.seek(0) - _buf = _io.BytesIO(_sf.read()) - with uproot.open(_buf, minimal_ttree_metadata=False) as _rf: - _num_file_branches = len(list(_rf[source].branches)) - if len(casc._branch_data) < _num_file_branches: + # check if file has TBranchElement branches (object dtype) + # _load_existing_ttree skips them, so we detect by checking dtype + if any( + bd.get("dtype") is not None and bd.get("dtype") == numpy.dtype("O") + for bd in casc._branch_data + ): raise NotImplementedError( "add_branches for files with TBranchElement branches is not yet " "supported via the cascade approach" @@ -2276,8 +2283,6 @@ def extend(self, data, *, accept_new_fields=False): for k, v in new_fields.items() } self.add_branches(zeros) - self._cascading.extend(self._file, self._file.sink, data) - return self._cascading.extend(self._file, self._file.sink, data) def show( From ae1f13506fcd95a63af73dbadcf2dd2e4de032b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:54:27 +0200 Subject: [PATCH 35/38] Use fIsRange to detect counter branches instead of name pattern matching --- src/uproot/writing/writable.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d4dafa8df..3dc7e087b 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1142,8 +1142,8 @@ def _get_chunk(start, stop): continue sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") # detect counter branches (e.g. njets for jagged jets array) - _branch_names = [br.name for br in branches] - _is_counter = b.name.startswith("n") and b.name[1:] in _branch_names + _leaves = b.member("fLeaves") + _is_counter = bool(_leaves) and bool(_leaves[0].member("fIsRange")) bd = { "fName": b.name, "branch_type": dtype, From 85a42810bc2380f1568bb28531c3dabf44e2299a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:56:45 +0200 Subject: [PATCH 36/38] Use key.cycle instead of hardcoded cycle number 1 in _load_existing_ttree --- src/uproot/writing/writable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 3dc7e087b..d05b29af6 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1241,7 +1241,7 @@ def _get_chunk(start, stop): ] } - dir_key = self._cascading.data.get_key(name, 1) + dir_key = self._cascading.data.get_key(name, key.cycle) freesegments = self._file._cascading.freesegments casc = ct.Tree.__new__(ct.Tree) From d0f2b8eda759e35441b1475c128fbb9d6c2a3900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:01:51 +0200 Subject: [PATCH 37/38] Fix subdirectory support in add_branches and _load_existing_ttree --- src/uproot/writing/writable.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d05b29af6..b930d2161 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1245,7 +1245,7 @@ def _get_chunk(start, stop): freesegments = self._file._cascading.freesegments casc = ct.Tree.__new__(ct.Tree) - casc._directory = self._file._cascading.rootdirectory + casc._directory = self._cascading casc._name = name casc._title = tree.title casc._freesegments = freesegments @@ -2116,9 +2116,14 @@ def add_branches(self, branches): source = self._path[-1] + # navigate to the correct directory (handles subdirectories) + directory = self._file.root_directory + for part in self._path[:-1]: + directory = directory[part] + # validate all branches have same length as existing tree - key = self._file._cascading.rootdirectory.data.get_key(source, 1) - casc = self._file.root_directory._load_existing_ttree(key)._cascading + key = directory._cascading.data.get_key(source) + casc = directory._load_existing_ttree(key)._cascading num_entries = casc._num_entries for branch_name, branch_data in branches.items(): @@ -2186,7 +2191,7 @@ def add_branches(self, branches): self._file.sink.flush() # update in-memory directory cache - dir_key_obj = self._file._cascading.rootdirectory.data.get_key(source, 1) + dir_key_obj = directory._cascading.data.get_key(source) dir_key_obj._seek_location = casc._key.seek_location # update self._cascading so subsequent extend uses correct metadata From d790b777f3403f95e55b643063a4a6ad13e8281b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jok=C5=ABbas=20Ma=C4=8Diulis?= <183698357+Yokubas@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:24:47 +0200 Subject: [PATCH 38/38] Fix counter branch validation in extend --- src/uproot/writing/writable.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index b930d2161..21db2ddf2 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2253,6 +2253,12 @@ def extend(self, data, *, accept_new_fields=False): if rn ) ] + # include counter branches that user explicitly provides in data + _counter_branch_names = [ + bd["fName"] + for bd in self._cascading._branch_data + if bd.get("kind") == "counter" and "fName" in bd + ] # check if data looks like a flat dict of branch arrays (not a record/awkward array) _data_is_flat_dict = isinstance(data, dict) and all( not hasattr(v, "fields") for v in data.values() @@ -2268,9 +2274,20 @@ def extend(self, data, *, accept_new_fields=False): new_fields = { k: v for k, v in data.items() - if k not in existing_names and k not in _record_parent_names + if k not in existing_names + and k not in _record_parent_names + and k not in _counter_branch_names } - missing = [b for b in existing_names if b not in data] + # skip counter branches not provided by user (auto-generated) + missing = [ + b + for b in existing_names + if b not in data and b not in _counter_branch_names + ] + # add counter branches to existing_names if user provides them + existing_names = _user_branch_names + [ + c for c in _counter_branch_names if c in data + ] if missing: raise ValueError( f"'extend' must fill every branch with the same number of entries; missing: {missing}"