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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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/55] 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}" From 4fa64d353b5650938ddeb42e7a9e135c216d1d94 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, 24 Aug 2026 12:19:10 +0300 Subject: [PATCH 39/55] fix: derive jagged branch dtype from interpretation content instead of numpy_dtype, preventing silent data corruption when extending an existing TTree via uproot.update() --- src/uproot/writing/writable.py | 18 +++++++++++++++++- tests/test_1690_ttree_inplace.py | 27 +++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 15327f732..f2577d341 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1153,10 +1153,26 @@ def _get_chunk(start, stop): for branch_idx, b in enumerate(branches): refs_list = list(b.cursor._refs.keys()) try: - dtype = b.interpretation.numpy_dtype.newbyteorder(">") + interpretation = b.interpretation + if isinstance(interpretation, uproot.interpretation.jagged.AsJagged): + # numpy_dtype of the AsJagged interpretation itself is + # dtype('O'); the basket actually holds the *content* + # dtype (e.g. float32 for "var * float32") + dtype = interpretation.content.to_dtype.newbyteorder(">") + else: + dtype = interpretation.numpy_dtype.newbyteorder(">") except AttributeError: # TBranchElement or other complex branch — skip continue + + if dtype.kind == "O": + raise NotImplementedError( + f"branch {b.name!r} has interpretation {interpretation!r}, " + "which is not yet supported by uproot.update(); only " + "numeric and jagged-numeric branches can currently be " + "extended on an existing TTree opened with uproot.update()" + ) + sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") # detect counter branches (e.g. njets for jagged jets array) _leaves = b.member("fLeaves") diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index b74fd6b4d..adb3dfd45 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -379,6 +379,33 @@ def test_extend_jagged_array(tmp_path): ] +def test_extend_jagged_array_new_session(tmp_path): + """Extending a jagged branch via uproot.update() must not corrupt data. + + Regression test: uproot.update() reconstructs branch metadata from disk + (_load_existing_ttree), which used to derive the on-disk dtype of a + jagged branch's content from the AsJagged interpretation's numpy_dtype + (always dtype('O')) instead of its content dtype, silently writing + garbage instead of raising. + """ + ak = pytest.importorskip("awkward") + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"jets": "var * float32"}) + f["tree"].extend({"jets": ak.Array([[1.0, 2.0], [3.0]])}) + + with uproot.update(path) as f: + f["tree"].extend({"jets": ak.Array([[7.0, 8.0, 9.0], [10.0]])}) + + with uproot.open(path) as f: + assert f["tree"]["jets"].array().tolist() == [ + [1.0, 2.0], + [3.0], + [7.0, 8.0, 9.0], + [10.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: From c8596a2a5e7a2a2dc52e4fcb28a3d5d39a3b28c2 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, 24 Aug 2026 12:25:47 +0300 Subject: [PATCH 40/55] fix: locate TTree/branch metadata offsets in _load_existing_ttree via content-independent byte markers instead of searching for zero-valued fields, preventing file corruption when extending a tree with zero baskets --- src/uproot/writing/writable.py | 65 +++++++++++++++++++++----------- tests/test_1690_ttree_inplace.py | 28 ++++++++++++++ 2 files changed, 71 insertions(+), 22 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index f2577d341..b22303711 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -474,6 +474,32 @@ def _move_tree(self, oldloc, newloc): self._trees[newloc] = tree +# the fixed 25-byte "empty TObjArray" that uproot.writing._cascadetree.Tree.write_anew +# writes twice per branch: once for the (always empty) TObjArray of sub-branches, and +# once for the TObjArray of embedded fBaskets, immediately before that branch's +# fBasketBytes/fBasketEntry/fBasketSeek arrays +_EMPTY_TOBJARRAY_BYTES = ( + b"@\x00\x00\x15\x00\x03\x00\x01\x00\x00\x00\x00\x03" + b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" +) + + +def _find_basket_metadata_start(raw, start): + """ + Locates the byte offset immediately after a branch's second + ``_EMPTY_TOBJARRAY_BYTES`` marker, i.e. where its + fBasketBytes/fBasketEntry/fBasketSeek arrays begin. + + This can't be found by searching for the contents of fBasketSeek[0] + (as a previous implementation did) because that value is 0 for any + branch that has never had a basket written, and a search for eight + zero bytes matches arbitrary unrelated data, corrupting the file. + """ + first = raw.find(_EMPTY_TOBJARRAY_BYTES, start) + second = raw.find(_EMPTY_TOBJARRAY_BYTES, first + len(_EMPTY_TOBJARRAY_BYTES)) + return second + len(_EMPTY_TOBJARRAY_BYTES) + + class WritableDirectory(MutableMapping): """ Args: @@ -1137,16 +1163,23 @@ def _get_chunk(start, stop): 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) + + # the fixed TAttLine v2 + TAttFill v2 + TAttMarker v2 block that + # uproot.writing._cascadetree.Tree.write_anew writes immediately before a + # TTree's fEntries/fTotBytes/fZipBytes metadata. Searching for this literal + # (rather than the *values* of fEntries/fTotBytes/fZipBytes, as a previous + # implementation did) is required because those values are all 0 for a + # freshly mktree'd tree with no entries, and a search for 24 zero bytes + # matches arbitrary unrelated data, corrupting the file. + _attline_attfill_attmarker = ( + b"@\x00\x00\x08\x00\x02\x02Z\x00\x01\x00\x01" + b"@\x00\x00\x06\x00\x02\x00\x00\x03\xe9" + b"@\x00\x00\n\x00\x02\x00\x01\x00\x01?\x80\x00\x00" ) - metadata_start = raw.find(seq) - if metadata_start == -1: + _attmarker_end = raw.find(_attline_attfill_attmarker) + if _attmarker_end == -1: raise RuntimeError(f"Could not find TTree metadata position in {name!r}") + metadata_start = _attmarker_end + len(_attline_attfill_attmarker) branch_data = [] branch_lookup = {} @@ -1211,20 +1244,8 @@ def _get_chunk(start, stop): ) - 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 - ) + "basket_metadata_start": _find_basket_metadata_start( + raw, b.cursor.index ), "tleaf_reference_number": ( refs_list[2 + branch_idx * 4] diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index adb3dfd45..777bfad81 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -187,6 +187,34 @@ def test_extend_preserves_existing(tmp_path): assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) +def test_extend_zero_basket_tree(tmp_path): + """Extending a freshly-mktree'd tree (no baskets written yet) via uproot.update(). + + Regression test: _load_existing_ttree used to locate the TTree's own + fEntries/fTotBytes/fZipBytes metadata, and each branch's + fBasketBytes/fBasketEntry/fBasketSeek arrays, by searching for the raw + bytes of their current (0, for a just-created tree) values. A search for + a run of zero bytes matches arbitrary unrelated data elsewhere in the + tree, corrupting the rewritten file instead of raising. + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + + with uproot.update(path) as f: + f["tree"].extend( + { + "x": np.arange(10, dtype=np.float32), + "y": np.arange(10, dtype=np.int32) * 2, + } + ) + + with uproot.open(path) as f: + assert f["tree"].num_entries == 10 + assert f["tree"]["x"].array().tolist() == list(range(10)) + assert f["tree"]["y"].array().tolist() == [i * 2 for i in range(10)] + + 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}) From 5907548ff00578bcba8b3affd0ec006556bae7c7 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, 24 Aug 2026 12:30:28 +0300 Subject: [PATCH 41/55] fix: support extending string branches in uproot.update() by recognizing AsStrings interpretation in _load_existing_ttree --- src/uproot/writing/writable.py | 11 +++++++++ tests/test_1690_ttree_inplace.py | 39 ++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index b22303711..39769344d 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1192,6 +1192,11 @@ def _get_chunk(start, stop): # dtype('O'); the basket actually holds the *content* # dtype (e.g. float32 for "var * float32") dtype = interpretation.content.to_dtype.newbyteorder(">") + elif isinstance( + interpretation, uproot.interpretation.strings.AsStrings + ): + # matches the dtype mktree/extend use for a "string" branch + dtype = numpy.dtype(str).newbyteorder(">") else: dtype = interpretation.numpy_dtype.newbyteorder(">") except AttributeError: @@ -1259,6 +1264,12 @@ def _get_chunk(start, stop): ), "tleaf_special_struct": _struct.Struct(">" + sc + sc), } + if dtype == ">U0": + # the length of the longest string written so far; extend() + # accumulates this as max(existing, new), so it must be + # seeded from what's already on disk or a later extend with + # only shorter strings would shrink fLen and corrupt reads + bd["fLen"] = b.member("fLeaves")[0].member("fLen") branch_data.append(bd) branch_lookup[b.name] = branch_idx diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 777bfad81..9279ce1bf 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -434,6 +434,45 @@ def test_extend_jagged_array_new_session(tmp_path): ] +def test_extend_string_branch_new_session(tmp_path): + """Extending a string branch via uproot.update() must work, not raise/corrupt. + + Regression test: _load_existing_ttree had no handling for the AsStrings + interpretation (numpy_dtype is dtype('O'), same as AsJagged), so + extending a string branch after reopening with uproot.update() failed. + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"s": "string"}) + f["tree"].extend({"s": ["a_very_long_string_here"]}) + + with uproot.update(path) as f: + f["tree"].extend({"s": ["x", "yy"]}) + + with uproot.open(path) as f: + assert f["tree"]["s"].array().tolist() == [ + "a_very_long_string_here", + "x", + "yy", + ] + + +def test_extend_string_branch_zero_basket(tmp_path): + """Extending a string branch that has never been extended (zero baskets).""" + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"s": "string", "x": np.float32}) + + with uproot.update(path) as f: + f["tree"].extend( + {"s": ["hi", "there"], "x": np.array([1.0, 2.0], dtype=np.float32)} + ) + + with uproot.open(path) as f: + assert f["tree"]["s"].array().tolist() == ["hi", "there"] + assert f["tree"]["x"].array().tolist() == [1.0, 2.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: From 255918904671b954a782c23cf34fcf8592e9bc5e 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, 24 Aug 2026 12:37:23 +0300 Subject: [PATCH 42/55] fix: robustly skip TBranchElement branches in _load_existing_ttree instead of crashing on access, and make the add_branches/extend TBranchElement guards actually reachable --- src/uproot/writing/_cascadetree.py | 1 + src/uproot/writing/writable.py | 39 +++++++++++++++++++++++------- tests/test_1690_ttree_inplace.py | 38 ++++++++++++++++++++++++++++- 3 files changed, 68 insertions(+), 10 deletions(-) diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index 205caa8ed..ba7b53b74 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -96,6 +96,7 @@ def __init__( self._field_name = field_name self._basket_capacity = initial_basket_capacity self._resize_factor = resize_factor + self._has_unsupported_branches = False if isinstance(branch_types, dict): branch_types_items = branch_types.items() diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 39769344d..de96a2d8f 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1184,6 +1184,16 @@ def _get_chunk(start, stop): branch_data = [] branch_lookup = {} for branch_idx, b in enumerate(branches): + if b.classname != "TBranch": + # TBranchElement, TBranchObject, etc. use an on-disk layout (split + # objects, embedded streamer info, header bytes in jagged baskets, + # ...) that this cascade — which only knows how to write plain + # TBranch/TLeaf — cannot reproduce. Leave them out of branch_data + # rather than crash trying to read TBranch-only fields (e.g. + # fMaximum, which TLeafElement doesn't have) or silently + # mis-write their baskets. + continue + refs_list = list(b.cursor._refs.keys()) try: interpretation = b.interpretation @@ -1271,7 +1281,9 @@ def _get_chunk(start, stop): # only shorter strings would shrink fLen and corrupt reads bd["fLen"] = b.member("fLeaves")[0].member("fLen") branch_data.append(bd) - branch_lookup[b.name] = branch_idx + # index into branch_data, not into the raw branches list: branches + # skipped above (non-TBranch) mean the two can diverge + branch_lookup[b.name] = len(branch_data) - 1 # fix counter references for jagged branches for bd in branch_data: @@ -1316,6 +1328,7 @@ def _get_chunk(start, stop): casc._freesegments = freesegments casc._branch_data = branch_data casc._branch_lookup = branch_lookup + casc._has_unsupported_branches = len(branch_data) != len(branches) casc._basket_capacity = ( next(iter(branches)).member("fMaxBaskets") if branches else 10 ) @@ -2201,15 +2214,14 @@ 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 (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 - ): + # _load_existing_ttree leaves TBranchElement (and other non-TBranch) + # branches out of casc._branch_data entirely, so rewriting the branch + # listing from casc._branch_data alone (as write_anew does) would + # silently drop them from the file + if casc._has_unsupported_branches: raise NotImplementedError( - "add_branches for files with TBranchElement branches is not yet " - "supported via the cascade approach" + "add_branches for files with TBranchElement (or other non-TBranch) " + "branches is not yet supported via the cascade approach" ) # add new branch dicts to cascade @@ -2299,6 +2311,15 @@ def extend(self, data, *, accept_new_fields=False): raise RuntimeError( "_cascading is None — this should not happen; please report this bug" ) + if self._cascading._has_unsupported_branches: + # _load_existing_ttree leaves TBranchElement (and other non-TBranch) + # branches out of _branch_data entirely, so extend() would only add + # entries to the branches it knows about, desynchronizing entry + # counts across the tree's branches + raise NotImplementedError( + "extend for files with TBranchElement (or other non-TBranch) " + "branches is not yet supported via the cascade approach" + ) # validate branches # get user-facing branch names (exclude auto-generated counter and record parent branches) # get record parent names to exclude their sub-fields diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 9279ce1bf..5fa4df7df 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -92,10 +92,46 @@ def test_add_branch_tbranchelement(tmp_path): ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - with pytest.raises(Exception): + with pytest.raises(NotImplementedError): f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) +def test_tbranchelement_access_does_not_crash(tmp_path): + """Merely accessing (not mutating) a TBranchElement tree under uproot.update(). + + Regression test: _load_existing_ttree used to include every branch whose + interpretation.numpy_dtype didn't raise AttributeError, including + TBranchElement branches with numeric-looking interpretations (e.g. + AsJagged content from split objects). Building that branch's metadata + then crashed reading TLeafElement.fMaximum, a member plain TLeaf has but + TLeafElement doesn't -- so simply doing f["events"] raised KeyInFileError. + """ + 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: + tree = f["events"] + assert tree.num_entries == 2421 + + +def test_extend_tbranchelement_raises(tmp_path): + """extend() on a TBranchElement file must raise, not silently desync entries. + + Regression test: _load_existing_ttree leaves unsupported (non-TBranch) + branches out of _branch_data, so extend() -- which only asks for the + branches it knows about -- would otherwise add entries to the supported + branches while leaving the TBranchElement branches' entry counts behind. + """ + 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): + f["events"].extend({"MC_leptonpdgid": np.zeros(1, dtype=np.int32)}) + + 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}) From d4d563efea1c3db31673f7eae18c05536d8e0175 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, 24 Aug 2026 13:00:46 +0300 Subject: [PATCH 43/55] fix: derive TTree/branch metadata offsets structurally instead of byte search, support fixed-size array branches, and reject divergent per-branch basket counts --- src/uproot/writing/_cascadetree.py | 24 ++++- src/uproot/writing/writable.py | 143 +++++++++++++++-------------- tests/test_1690_ttree_inplace.py | 65 ++++++++++++- 3 files changed, 157 insertions(+), 75 deletions(-) diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index ba7b53b74..c614c3813 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -97,6 +97,7 @@ def __init__( self._basket_capacity = initial_basket_capacity self._resize_factor = resize_factor self._has_unsupported_branches = False + self._has_divergent_baskets = False if isinstance(branch_types, dict): branch_types_items = branch_types.items() @@ -823,7 +824,24 @@ def extend(self, file, sink, data): self.write_updates(sink) - def write_anew(self, sink): + def _build_out(self): + """ + Serializes this TTree's metadata blob (TTree + TBranches + TLeaves, + everything but the TBaskets themselves) into a list of byte chunks, + exactly as ``write_anew`` writes it. + + As a side effect, this sets ``self._metadata_start`` and, for every + branch, ``datum["metadata_start"]``, ``datum["basket_metadata_start"]``, + and ``datum["tleaf_reference_number"]`` — the byte offsets (relative to + this blob) that ``write_updates`` writes at directly, without going + through ``write_anew`` again. This is the single source of truth for + that layout: computing it any other way (e.g. searching a decompressed + blob for the current value of a field) risks disagreeing with it, + especially when that field's current value is 0 and matches unrelated + bytes. ``uproot.writing.writable.WritableDirectory._load_existing_ttree`` + calls this once, on a freshly reconstructed cascade, purely to recover + these offsets — the returned ``out`` is discarded, nothing is written. + """ key_num_bytes = uproot.reading._key_format_big.size + 6 name_asbytes = self._name.encode(errors="surrogateescape") title_asbytes = self._title.encode(errors="surrogateescape") @@ -1201,6 +1219,10 @@ def write_anew(self, sink): self._metadata_start = sum(len(x) for x in out[:metadata_out_index]) + return out + + def write_anew(self, sink): + out = self._build_out() raw_data = b"".join(out) self._key = self._directory.add_object( sink, diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index de96a2d8f..d43560535 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -474,32 +474,6 @@ def _move_tree(self, oldloc, newloc): self._trees[newloc] = tree -# the fixed 25-byte "empty TObjArray" that uproot.writing._cascadetree.Tree.write_anew -# writes twice per branch: once for the (always empty) TObjArray of sub-branches, and -# once for the TObjArray of embedded fBaskets, immediately before that branch's -# fBasketBytes/fBasketEntry/fBasketSeek arrays -_EMPTY_TOBJARRAY_BYTES = ( - b"@\x00\x00\x15\x00\x03\x00\x01\x00\x00\x00\x00\x03" - b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" -) - - -def _find_basket_metadata_start(raw, start): - """ - Locates the byte offset immediately after a branch's second - ``_EMPTY_TOBJARRAY_BYTES`` marker, i.e. where its - fBasketBytes/fBasketEntry/fBasketSeek arrays begin. - - This can't be found by searching for the contents of fBasketSeek[0] - (as a previous implementation did) because that value is 0 for any - branch that has never had a basket written, and a search for eight - zero bytes matches arbitrary unrelated data, corrupting the file. - """ - first = raw.find(_EMPTY_TOBJARRAY_BYTES, start) - second = raw.find(_EMPTY_TOBJARRAY_BYTES, first + len(_EMPTY_TOBJARRAY_BYTES)) - return second + len(_EMPTY_TOBJARRAY_BYTES) - - class WritableDirectory(MutableMapping): """ Args: @@ -1159,31 +1133,20 @@ def _get_chunk(start, stop): ) 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") - # the fixed TAttLine v2 + TAttFill v2 + TAttMarker v2 block that - # uproot.writing._cascadetree.Tree.write_anew writes immediately before a - # TTree's fEntries/fTotBytes/fZipBytes metadata. Searching for this literal - # (rather than the *values* of fEntries/fTotBytes/fZipBytes, as a previous - # implementation did) is required because those values are all 0 for a - # freshly mktree'd tree with no entries, and a search for 24 zero bytes - # matches arbitrary unrelated data, corrupting the file. - _attline_attfill_attmarker = ( - b"@\x00\x00\x08\x00\x02\x02Z\x00\x01\x00\x01" - b"@\x00\x00\x06\x00\x02\x00\x00\x03\xe9" - b"@\x00\x00\n\x00\x02\x00\x01\x00\x01?\x80\x00\x00" - ) - _attmarker_end = raw.find(_attline_attfill_attmarker) - if _attmarker_end == -1: - raise RuntimeError(f"Could not find TTree metadata position in {name!r}") - metadata_start = _attmarker_end + len(_attline_attfill_attmarker) + # metadata_start (TTree-level) and, per branch, metadata_start and + # basket_metadata_start are filled in below by casc._build_out(), which + # walks this cascade's layout the same way write_anew does -- rather than + # recovering them by searching the decompressed blob for the *current + # value* of a nearby field, which breaks whenever that value is 0 (e.g. a + # freshly mktree'd tree has fEntries == fTotBytes == fZipBytes == 0, and a + # search for zero bytes matches arbitrary unrelated data). branch_data = [] branch_lookup = {} - for branch_idx, b in enumerate(branches): + for b in branches: if b.classname != "TBranch": # TBranchElement, TBranchObject, etc. use an on-disk layout (split # objects, embedded streamer info, header bytes in jagged baskets, @@ -1194,7 +1157,6 @@ def _get_chunk(start, stop): # mis-write their baskets. continue - refs_list = list(b.cursor._refs.keys()) try: interpretation = b.interpretation if isinstance(interpretation, uproot.interpretation.jagged.AsJagged): @@ -1221,6 +1183,13 @@ def _get_chunk(start, stop): "extended on an existing TTree opened with uproot.update()" ) + # a fixed-size array branch (e.g. "float[3]") has a numpy_dtype with + # a subdtype/shape, like _branch_np splits for a freshly created branch + if dtype.subdtype is not None: + dtype, shape = dtype.subdtype + else: + shape = () + sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") # detect counter branches (e.g. njets for jagged jets array) _leaves = b.member("fLeaves") @@ -1231,7 +1200,7 @@ def _get_chunk(start, stop): "kind": "counter" if _is_counter else "normal", "counter": None, "dtype": dtype, - "shape": (), + "shape": shape, "fTitle": b.member("fTitle"), "compression": b.compression, "fBasketSize": b.member("fBasketSize"), @@ -1246,27 +1215,10 @@ def _get_chunk(start, stop): "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": _find_basket_metadata_start( - raw, b.cursor.index - ), - "tleaf_reference_number": ( - refs_list[2 + branch_idx * 4] - if 2 + branch_idx * 4 < len(refs_list) - else 0 - ), + # filled in below by casc._build_out() + "metadata_start": None, + "basket_metadata_start": None, + "tleaf_reference_number": 0, "tleaf_maximum_value": ( int(b.member("fLeaves")[0].member("fMaximum")) if b.member("fLeaves") @@ -1295,7 +1247,23 @@ def _get_chunk(start, stop): if counter_bd is not None: bd["counter"] = counter_bd - fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 + # fWriteBasket/fMaxBaskets are only tracked once per tree (not per + # branch) by this cascade, but nothing guarantees every branch in a + # preexisting tree agrees on either -- ROOT commonly flushes a basket + # once a branch's accumulated data exceeds fBasketSize, so branches + # with different per-entry sizes accumulate baskets at different + # rates even when filled together from the start. Detect that here + # so extend()/add_branches() can refuse rather than silently use one + # branch's basket bookkeeping for all of them. + _supported_branches = [b for b in branches if b.classname == "TBranch"] + _write_basket_values = {b.member("fWriteBasket") for b in _supported_branches} + _max_basket_values = {b.member("fMaxBaskets") for b in _supported_branches} + has_divergent_baskets = ( + len(_write_basket_values) > 1 or len(_max_basket_values) > 1 + ) + fWriteBasket = ( + _supported_branches[0].member("fWriteBasket") if _supported_branches else 0 + ) metadata = { k: tree.member(k) for k in [ @@ -1329,18 +1297,25 @@ def _get_chunk(start, stop): casc._branch_data = branch_data casc._branch_lookup = branch_lookup casc._has_unsupported_branches = len(branch_data) != len(branches) + casc._has_divergent_baskets = has_divergent_baskets casc._basket_capacity = ( - next(iter(branches)).member("fMaxBaskets") if branches else 10 + _supported_branches[0].member("fMaxBaskets") if _supported_branches else 10 ) casc._resize_factor = 10.0 casc._counter_name = lambda counted: "n" + counted casc._field_name = None - casc._metadata_start = metadata_start casc._num_baskets = fWriteBasket casc._num_entries = fEntries casc._metadata = metadata casc._key = dir_key + # recovers casc._metadata_start and each branch's metadata_start / + # basket_metadata_start / tleaf_reference_number by walking the same + # layout write_anew would emit for this tree, rather than searching the + # blob for byte patterns; the returned chunks are discarded, nothing is + # written to the file + casc._build_out() + path = (*self._path, name) writable_tree = WritableTree(path, self._file, casc) self._file._trees[key.seek_location] = writable_tree @@ -2223,6 +2198,18 @@ def add_branches(self, branches): "add_branches for files with TBranchElement (or other non-TBranch) " "branches is not yet supported via the cascade approach" ) + if casc._has_divergent_baskets: + # this cascade tracks one fWriteBasket/fMaxBaskets pair per tree, not + # per branch, but nothing guarantees every branch in a preexisting + # tree agrees on either (e.g. a ROOT-written tree whose branches + # have different per-entry sizes, and so flush baskets at different + # rates); rewriting the branch listing would apply one branch's + # basket bookkeeping to every branch, corrupting the others + raise NotImplementedError( + "add_branches for a TTree whose branches do not all have the same " + "number of baskets / basket capacity is not yet supported via the " + "cascade approach" + ) # add new branch dicts to cascade compression = casc._freesegments.fileheader.compression @@ -2320,6 +2307,20 @@ def extend(self, data, *, accept_new_fields=False): "extend for files with TBranchElement (or other non-TBranch) " "branches is not yet supported via the cascade approach" ) + if self._cascading._has_divergent_baskets: + # this cascade tracks one fWriteBasket/fMaxBaskets pair per tree, not + # per branch, but nothing guarantees every branch in a preexisting + # tree agrees on either (e.g. a ROOT-written tree whose branches + # have different per-entry sizes, and so flush baskets at different + # rates); extend() would write every branch's new basket at an + # offset computed from one branch's basket count, corrupting the + # others (or crash outright if that count exceeds another + # branch's actual basket capacity) + raise NotImplementedError( + "extend for a TTree whose branches do not all have the same " + "number of baskets / basket capacity is not yet supported via " + "the cascade approach" + ) # validate branches # get user-facing branch names (exclude auto-generated counter and record parent branches) # get record parent names to exclude their sub-fields diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 5fa4df7df..6800d4572 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -85,8 +85,10 @@ 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 + # add_branches for TBranchElement files is not supported: _load_existing_ttree + # leaves TBranchElement branches out of branch_data (it can only write plain + # TBranch), so rewriting the branch listing from branch_data would silently + # drop them shutil.copy( data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") ) @@ -178,7 +180,7 @@ def test_add_branch_tbranchelement_root_readable(tmp_path): ) with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: - with pytest.raises((NotImplementedError, TypeError, KeyError)): + with pytest.raises(NotImplementedError): f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) @@ -509,6 +511,63 @@ def test_extend_string_branch_zero_basket(tmp_path): assert f["tree"]["x"].array().tolist() == [1.0, 2.0] +def test_access_fixed_size_array_branch(tmp_path): + """Accessing a real ROOT-written tree with fixed-size array branches (e.g. "bool[3]"). + + Regression test: _load_existing_ttree hardcoded every branch's "shape" to + () and used interpretation.numpy_dtype directly. For a fixed-size array + branch, that dtype carries a subdtype/shape (e.g. dtype(('?', (3,))) for + "bool[3]"), which isn't a key in _dtype_to_char, so even plain access (not + just extend) crashed with a KeyError. + """ + path = os.path.join(tmp_path, "sample.root") + shutil.copy(data_path("uproot-sample-6.20.04-uncompressed.root"), path) + + with uproot.update(path) as f: + assert f["sample"].num_entries == 30 + + +@skip_no_root +def test_extend_divergent_basket_counts_raises(tmp_path): + """extend() on a ROOT-written tree whose branches have different basket counts. + + Regression test: this cascade tracks one fWriteBasket/fMaxBaskets pair per + tree (taken from a single branch), not per branch. ROOT commonly flushes a + basket once a branch's accumulated data exceeds fBasketSize, so branches + with different per-entry sizes accumulate baskets at different rates even + when filled together from the start -- applying one branch's basket count + to every branch corrupted or crashed the file. It must now raise instead. + """ + import array + + path = os.path.join(tmp_path, "divergent.root") + rf = ROOT.TFile(str(path), "RECREATE") + rt = ROOT.TTree("tree", "tree") + x = array.array("f", [0.0]) + y = array.array("d", [0.0]) + # small basket size + different per-entry byte sizes (4 vs 8 bytes) so the + # two branches flush baskets at different rates + rt.Branch("x", x, "x/F", 64) + rt.Branch("y", y, "y/D", 64) + for i in range(100): + x[0] = float(i) + y[0] = float(i) * 2 + rt.Fill() + rt.Write() + rf.Close() + + with uproot.update(path) as f: + tree = f["tree"] + assert tree._cascading._has_divergent_baskets + with pytest.raises(NotImplementedError): + tree.extend( + { + "x": np.array([999.0], dtype=np.float32), + "y": np.array([888.0], dtype=np.float64), + } + ) + + 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: From ae8ee2fa963fceadbf11f1db0d9a57894bcb8322 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, 25 Aug 2026 14:41:43 +0300 Subject: [PATCH 44/55] fix: correct fWriteBasket for brand-new branches in add_branches instead of stamping them with the whole tree's basket count --- src/uproot/writing/writable.py | 23 ++++++++++++++++ tests/test_1690_ttree_inplace.py | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d43560535..01a4f1fe0 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2252,6 +2252,29 @@ def add_branches(self, branches): casc._num_baskets = old_num_baskets casc.write_updates(self._file.sink) + + # write_updates() stamps every branch's fWriteBasket with the tree-wide + # casc._num_baskets (old_num_baskets) -- correct for the pre-existing + # branches, which really do have that many baskets, but wrong for the + # brand-new branches just written above: add_branches always writes + # exactly one basket per new branch (indexed 0 above), so their real + # fWriteBasket is 1, not old_num_baskets. Left uncorrected, the file + # claims baskets the new branch never wrote, and reading it back later + # fails to account for the true number of entries. + base = casc._key.seek_location + casc._key.num_bytes + for branch_name in branches: + datum = casc._branch_data[casc._branch_lookup[branch_name]] + self._file.sink.write( + base + datum["metadata_start"], + uproot.models.TBranch._tbranch13_format1.pack( + 0, # fCompress (write_updates also always writes 0 here; see its comment) + datum["fBasketSize"], + datum["fEntryOffsetLen"], + 1, # fWriteBasket -- this branch has exactly one basket + casc._num_entries, # fEntryNumber + ), + ) + self._file.sink.flush() # update in-memory directory cache diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 6800d4572..6a7a6ef6a 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -362,6 +362,53 @@ def test_add_branch_sequential(tmp_path): assert np.all(f["tree"]["branch_b"].array() == 3) +def test_add_branch_after_multiple_extends(tmp_path): + """add_branches() on a tree that already has more than one basket. + + Regression test: write_updates() stamps every branch's fWriteBasket with + the tree-wide casc._num_baskets, correct for the pre-existing branches + (which really do have that many baskets) but wrong for a brand-new + branch, which add_branches() always writes exactly one basket for. Left + uncorrected, the new branch's fWriteBasket claimed as many baskets as the + rest of the tree even though only basket 0 held real data. That alone + doesn't crash a plain read (this file's new_branch reads back correctly + below), but it makes fWriteBasket agree with the older branches' basket + count despite the real per-branch layout being divergent -- which let a + follow-up extend() sail past the divergent-basket-count guard (since + that guard trusts fWriteBasket) and write another basket for new_branch + indexed as if it were basket 2, when only basket 0 was ever real. Reading + it back then failed with a ValueError about basket/entry counts not + adding up. Every existing add_branches test happened to extend() exactly + once first, so old_num_baskets was coincidentally always 1 and this + never surfaced. With fWriteBasket corrected, the guard now sees the true + divergence and rejects the follow-up extend() cleanly instead. + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.full(100, 1.0, dtype=np.float32)}) + f["tree"].extend({"x": np.full(100, 2.0, dtype=np.float32)}) + + with uproot.update(path) as f: + f["tree"].add_branches({"new_branch": np.full(200, 9.0, dtype=np.float32)}) + + with uproot.open(path) as f: + assert f["tree"].num_entries == 200 + assert np.all(f["tree"]["new_branch"].array() == 9.0) + x = f["tree"]["x"].array() + assert np.all(x[:100] == 1.0) + assert np.all(x[100:] == 2.0) + + with uproot.update(path) as f: + with pytest.raises(NotImplementedError): + f["tree"].extend( + { + "x": np.full(50, 3.0, dtype=np.float32), + "new_branch": np.full(50, 8.0, dtype=np.float32), + } + ) + + 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}) From 872f9a3fa16ad9238ceb0866513be7567a6f200c 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, 25 Aug 2026 14:46:26 +0300 Subject: [PATCH 45/55] docs: correct add_branches docstring claiming TBranchElement support it doesn't have --- src/uproot/writing/writable.py | 21 ++++++++++++++++++++- tests/test_1690_ttree_inplace.py | 12 ++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 01a4f1fe0..236cc6e46 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2157,12 +2157,20 @@ def add_branches(self, branches): 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["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + .. note:: + + Only trees whose branches are all plain ``TBranch`` (not + ``TBranchElement``, e.g. split objects) are supported. Rewriting + the branch listing for a tree with ``TBranchElement`` branches + would silently drop them, so this raises ``NotImplementedError`` + for such trees instead. This mirrors the same restriction on + :ref:`uproot.writing.writable.WritableTree.extend`. """ if self._file.sink.closed: raise ValueError("cannot modify a TTree in a closed file") @@ -2316,6 +2324,17 @@ def extend(self, data, *, accept_new_fields=False): .. warning:: **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) `__. + + .. note:: + + On a tree reopened with :doc:`uproot.writing.writable.WritableDirectory.update`, + only trees whose branches are all plain ``TBranch`` (not + ``TBranchElement``, e.g. split objects) are supported, and only + trees whose branches all agree on basket count/capacity (true by + construction for a tree Uproot itself wrote and has not touched + with :ref:`uproot.writing.writable.WritableTree.add_branches` + since). Either case raises ``NotImplementedError`` rather than + desynchronizing or corrupting the tree's branches. """ if self._cascading is None: raise RuntimeError( diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 6a7a6ef6a..f34583c89 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -98,6 +98,18 @@ def test_add_branch_tbranchelement(tmp_path): f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) +def test_add_branches_docstring_does_not_claim_tbranchelement_support(): + """add_branches()'s docstring must not claim TBranchElement support it doesn't have. + + Regression test: the docstring said "Works with both simple TBranch and + TBranchElement files," directly contradicted by the NotImplementedError + add_branches() raises for exactly that case (see + test_add_branch_tbranchelement above). + """ + doc = uproot.writing.writable.WritableTree.add_branches.__doc__ + assert "Works with both simple TBranch and TBranchElement files" not in doc + + def test_tbranchelement_access_does_not_crash(tmp_path): """Merely accessing (not mutating) a TBranchElement tree under uproot.update(). From a1dd1dffe087470657241824dd326cc0a16268ec 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, 25 Aug 2026 14:51:10 +0300 Subject: [PATCH 46/55] fix: don't infer a counter branch for a string branch's coincidentally-named n sibling --- src/uproot/writing/writable.py | 15 ++++++++++++++- tests/test_1690_ttree_inplace.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 236cc6e46..bf1e7cf4b 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1238,8 +1238,21 @@ def _get_chunk(start, stop): branch_lookup[b.name] = len(branch_data) - 1 # fix counter references for jagged branches + # + # fEntryOffsetLen > 0 alone isn't a reliable signal that a branch is + # jagged-with-a-counter: a string branch also gets fEntryOffsetLen > 0 + # once it has data (its basket-internal offset table), but strings are + # never counted by a separate branch. Without the dtype check, a string + # branch could be wrongly paired with an unrelated same-named "n"+name + # branch (e.g. string branch "id" and an unrelated int branch "nid"), + # making extend()'s counter/no-counter branch treat it as jagged + # numeric instead of a string. for bd in branch_data: - if bd.get("fEntryOffsetLen", 0) > 0 and bd["counter"] is None: + if ( + bd.get("fEntryOffsetLen", 0) > 0 + and bd["counter"] is None + and bd.get("dtype") != ">U0" + ): counter_nm = "n" + bd["fName"] counter_bd = next( (x for x in branch_data if x["fName"] == counter_nm), None diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index f34583c89..a7993ec52 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -570,6 +570,37 @@ def test_extend_string_branch_zero_basket(tmp_path): assert f["tree"]["x"].array().tolist() == [1.0, 2.0] +def test_extend_string_branch_with_coincidentally_named_counter(tmp_path): + """A string branch must not be misattributed a counter from an unrelated same-named branch. + + Regression test: the counter-branch inference in _load_existing_ttree + matched any branch with fEntryOffsetLen > 0 to a same-named "n"+branch, + without checking the branch was actually jagged-content (numeric). A + string branch also gets fEntryOffsetLen > 0 once it has data (its + basket-internal offset table), so a string branch named "id" alongside + an unrelated int branch named "nid" got "nid" wrongly attached as its + counter -- making extend() treat "id" via the jagged/counted code path + (which expects an Awkward array with a .layout) instead of the string + path, raising a confusing AttributeError on the very next extend(). + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"id": "string", "nid": np.int32}) + f["tree"].extend( + {"id": ["aaa", "bb"], "nid": np.array([10, 20], dtype=np.int32)} + ) + + with uproot.update(path) as f: + t = f["tree"] + id_bd = next(bd for bd in t._cascading._branch_data if bd["fName"] == "id") + assert id_bd["counter"] is None + t.extend({"id": ["ccc"], "nid": np.array([30], dtype=np.int32)}) + + with uproot.open(path) as f: + assert f["tree"]["id"].array().tolist() == ["aaa", "bb", "ccc"] + assert f["tree"]["nid"].array().tolist() == [10, 20, 30] + + def test_access_fixed_size_array_branch(tmp_path): """Accessing a real ROOT-written tree with fixed-size array branches (e.g. "bool[3]"). From 5a4ace0fcf0db8f688513a14b9edc29201c75b99 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, 25 Aug 2026 14:55:42 +0300 Subject: [PATCH 47/55] fix: move the WritableFile._trees cache entry when add_branches relocates the tree, instead of leaking a stale one --- src/uproot/writing/writable.py | 9 ++++++++- tests/test_1690_ttree_inplace.py | 27 +++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index bf1e7cf4b..37359ce5d 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2247,8 +2247,15 @@ def add_branches(self, branches): casc._branch_data.append(new_bd) casc._branch_lookup[branch_name] = len(casc._branch_data) - 1 - # rewrite TTree metadata blob with new branches included + # rewrite TTree metadata blob with new branches included -- this relocates + # the tree (frees its old space, allocates new space for the larger + # blob), so the WritableFile._trees cache needs to move with it, the same + # way extend()'s own basket-capacity-expansion relocation does via + # file._move_tree(); otherwise a stale entry at the old location lingers + # in that cache indefinitely + oldloc = casc._key.seek_location casc.write_anew(self._file.sink) + self._file._move_tree(oldloc, casc._key.seek_location) # write one basket per new branch old_num_baskets = casc._num_baskets diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index a7993ec52..adafbe183 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -374,6 +374,33 @@ def test_add_branch_sequential(tmp_path): assert np.all(f["tree"]["branch_b"].array() == 3) +def test_add_branch_does_not_leak_stale_trees_cache_entries(tmp_path): + """add_branches() must not leave stale entries in WritableFile._trees behind. + + Regression test: add_branches() calls write_anew(), which relocates the + tree (frees its old space, allocates new space for the larger blob), the + same way extend()'s basket-capacity-expansion path does -- but unlike + that path, add_branches() never called file._move_tree() to move the + WritableFile._trees cache entry to the new location, so the entry at the + old (now-freed) location was never cleaned up. Every add_branches() call + in a session left one more stale entry behind. + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(path) as f: + before = len(f._file._trees) + f["tree"].add_branches({"a": np.ones(100, dtype=np.float32)}) + f["tree"].add_branches({"b": np.ones(100, dtype=np.float32)}) + f["tree"].add_branches({"c": np.ones(100, dtype=np.float32)}) + assert len(f._file._trees) == before + 1 + + with uproot.open(path) as f: + assert f["tree"].arrays().fields == ["x", "a", "b", "c"] + + def test_add_branch_after_multiple_extends(tmp_path): """add_branches() on a tree that already has more than one basket. From 9dbba9861358a9c80b110ffa0d6b27872c6db144 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, 25 Aug 2026 14:59:39 +0300 Subject: [PATCH 48/55] perf: reuse self._cascading in add_branches instead of unconditionally reloading and re-parsing the tree from disk --- src/uproot/writing/writable.py | 13 ++++++++++--- tests/test_1690_ttree_inplace.py | 27 +++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 37359ce5d..6ae0eaf77 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2195,9 +2195,16 @@ def add_branches(self, branches): for part in self._path[:-1]: directory = directory[part] - # validate all branches have same length as existing tree - key = directory._cascading.data.get_key(source) - casc = directory._load_existing_ttree(key)._cascading + # validate all branches have same length as existing tree; self._cascading + # is already fully flushed and current on disk by the end of every + # extend()/add_branches() call (the same trust extend() itself places in + # it), so there's no need to re-read and re-parse the whole tree from + # disk here + if self._cascading is None: + raise RuntimeError( + "_cascading is None — this should not happen; please report this bug" + ) + casc = self._cascading num_entries = casc._num_entries for branch_name, branch_data in branches.items(): diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index adafbe183..43d98329a 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -401,6 +401,33 @@ def test_add_branch_does_not_leak_stale_trees_cache_entries(tmp_path): assert f["tree"].arrays().fields == ["x", "a", "b", "c"] +def test_add_branch_does_not_reload_tree_from_disk(tmp_path): + """add_branches() should reuse self._cascading, not re-read the tree from disk. + + Efficiency regression test: add_branches() used to unconditionally call + _load_existing_ttree() again, re-reading and re-parsing every branch's + metadata from disk even though self._cascading was already current + (extend() itself already trusts self._cascading without reloading it). + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(path) as f: + tree = f["tree"] + calls = [] + original = type(f)._load_existing_ttree + type(f)._load_existing_ttree = lambda self, key: ( + calls.append(1) or original(self, key) + ) + try: + tree.add_branches({"a": np.ones(100, dtype=np.float32)}) + finally: + type(f)._load_existing_ttree = original + assert calls == [] + + def test_add_branch_after_multiple_extends(tmp_path): """add_branches() on a tree that already has more than one basket. From a14d8b0e35fa5ea05ebfe6e4aa6be9c5500abc0c 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, 26 Aug 2026 15:32:57 +0300 Subject: [PATCH 49/55] fix: force one full write_anew before extend's first incremental patch on a preexisting tree, fixing silent corruption when extending a ROOT-written TTree --- src/uproot/writing/_cascadetree.py | 61 ++++++++++++++++++++++++------ src/uproot/writing/writable.py | 14 +++++++ tests/test_1690_ttree_inplace.py | 56 +++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 11 deletions(-) diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index c614c3813..158808da8 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -98,6 +98,7 @@ def __init__( self._resize_factor = resize_factor self._has_unsupported_branches = False self._has_divergent_baskets = False + self._needs_relocation_before_extend = False if isinstance(branch_types, dict): branch_types_items = branch_types.items() @@ -448,7 +449,55 @@ def num_entries(self): def num_baskets(self): return self._num_baskets + def _relocate(self, file, sink): + """ + Rewrites this tree's whole metadata blob from scratch via write_anew(), + relocating it (freeing its old space, allocating new space), and moves + the WritableFile._trees cache entry to match. Shared by the two + situations that need a full rewrite rather than an incremental + write_updates() patch: basket-capacity growth, and (see + _needs_relocation_before_extend) establishing Uproot's own canonical + layout for a tree write_updates() hasn't verified it can trust yet. + """ + # seek_location is the tree's actual physical on-disk position (it's + # what every sink.read(key.seek_location, ...) call in + # _load_existing_ttree uses), and what file._move_tree()/ + # WritableFile._trees are keyed by everywhere else in the codebase. + # location coincides with it for any tree Uproot itself has written + # (all our own tests, hence never catching this), but can diverge for + # a preexisting (e.g. ROOT-written) tree's Key object -- silently + # freeing/keying the wrong byte range and cache entry. + oldloc = start = self._key.seek_location + stop = start + self._key.num_bytes + self._key.compressed_bytes + + self.write_anew(sink) + + newloc = self._key.seek_location + file._move_tree(oldloc, newloc) + + self._freesegments.release(start, stop) + sink.set_file_length(self._freesegments.fileheader.end) + sink.flush() + def extend(self, file, sink, data): + if self._needs_relocation_before_extend: + # metadata_start/basket_metadata_start (computed structurally by + # _build_out(), see _load_existing_ttree) are only guaranteed + # correct for a tree actually laid out the way Uproot's own + # write_anew() lays one out. A tree loaded from an existing file + # may have been written by ROOT instead, with a byte-for-byte + # different (but semantically equivalent) layout Uproot's writer + # doesn't replicate -- write_updates() patching at Uproot's + # assumed offsets would then silently corrupt unrelated bytes. + # write_anew() here establishes Uproot's own canonical layout for + # real (the same thing add_branches() already always does, which + # is why it already works on ROOT-written files), so every + # write_updates() patch from here on is guaranteed to target the + # right bytes. Costs one full blob rewrite, but only once, on the + # first extend() after loading -- not on every call. + self._relocate(file, sink) + self._needs_relocation_before_extend = False + # expand capacity if this would REACH (not EXCEED) the existing capacity # that's because completely a full fBasketEntry has nowhere to put the # number of entries in the last basket (it's a fencepost principle thing), @@ -481,17 +530,7 @@ def extend(self, file, sink, data): datum["fBasketSeek"][: len(fBasketSeek)] = fBasketSeek datum["fBasketEntry"][len(fBasketEntry)] = self._num_entries - oldloc = start = self._key.location - stop = start + self._key.num_bytes + self._key.compressed_bytes - - self.write_anew(sink) - - newloc = self._key.seek_location - file._move_tree(oldloc, newloc) - - self._freesegments.release(start, stop) - sink.set_file_length(self._freesegments.fileheader.end) - sink.flush() + self._relocate(file, sink) provided = None diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 6ae0eaf77..30749c187 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1311,6 +1311,16 @@ def _get_chunk(start, stop): casc._branch_lookup = branch_lookup casc._has_unsupported_branches = len(branch_data) != len(branches) casc._has_divergent_baskets = has_divergent_baskets + # metadata_start/basket_metadata_start were just derived structurally, + # assuming this tree is laid out the way Uproot's own writer lays one + # out. That's only actually true once Uproot itself has written this + # tree at least once; a preexisting tree may be ROOT-written, with a + # different byte layout write_updates() would patch at the wrong + # offsets. extend() checks this and does one full write_anew() before + # its first incremental patch to establish Uproot's own layout for + # real. add_branches() already always calls write_anew() regardless, + # so it clears this flag itself once it does. + casc._needs_relocation_before_extend = True casc._basket_capacity = ( _supported_branches[0].member("fMaxBaskets") if _supported_branches else 10 ) @@ -2263,6 +2273,10 @@ def add_branches(self, branches): oldloc = casc._key.seek_location casc.write_anew(self._file.sink) self._file._move_tree(oldloc, casc._key.seek_location) + # this write_anew() just established Uproot's own canonical layout for + # real, so a subsequent extend() on this same cascade doesn't need to + # redundantly do it again + casc._needs_relocation_before_extend = False # write one basket per new branch old_num_baskets = casc._num_baskets diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 43d98329a..a31e0d3ba 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -475,6 +475,62 @@ def test_add_branch_after_multiple_extends(tmp_path): ) +def test_extend_root_written_tree(tmp_path): + """extend() on a genuinely ROOT-written tree must not corrupt the file. + + Regression test: metadata_start/basket_metadata_start are derived + structurally (see _build_out()), which assumes the tree is laid out the + way Uproot's own writer lays one out. A ROOT-written tree can have a + byte-for-byte different (but semantically equivalent) layout, so + write_updates() -- which extend() uses for every call except capacity + growth -- patched the wrong bytes with no exception raised at write + time, corrupting an existing basket's compressed data badly enough that + reading it back raised a zlib decompression error. add_branches() + already always calls write_anew() and so never hit this; extend() now + does the same once, on the first call after loading a preexisting tree, + establishing Uproot's own canonical layout before ever trusting a + write_updates() patch. + """ + path = os.path.join(tmp_path, "test.root") + shutil.copy(data_path("uproot-foriter.root"), path) + + with uproot.open(path) as f: + before = f["foriter"]["data"].array().tolist() + + with uproot.update(path) as f: + f["foriter"].extend({"data": np.arange(2, dtype=np.int32)}) + + with uproot.open(path) as f: + after = f["foriter"]["data"].array().tolist() + assert after == before + [0, 1] + + +def test_extend_root_written_tree_multiple_branches(tmp_path): + """extend() on a genuinely ROOT-written tree with many branches, twice in a row.""" + path = os.path.join(tmp_path, "test.root") + shutil.copy(data_path("uproot-Zmumu.root"), path) + + with uproot.open(path) as f: + before = f["events"].arrays() + + with uproot.update(path) as f: + data = {name: np.asarray(before[name][:3]) for name in before.fields} + f["events"].extend(data) + + with uproot.update(path) as f: + data = {name: np.asarray(before[name][3:5]) for name in before.fields} + f["events"].extend(data) + + with uproot.open(path) as f: + after = f["events"].arrays() + assert len(after) == len(before) + 5 + assert np.allclose(after["px1"][: len(before)], before["px1"]) + assert np.allclose( + after["px1"][len(before) : len(before) + 3], before["px1"][:3] + ) + assert np.allclose(after["px1"][len(before) + 3 :], before["px1"][3:5]) + + 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}) From 79dde8333413e31d1d0dae2ee5e103e82706988a 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, 26 Aug 2026 17:10:15 +0300 Subject: [PATCH 50/55] fix: refuse add_branches up front on a tree that already has more than one basket --- src/uproot/writing/writable.py | 16 +++++++++++ tests/test_1690_ttree_inplace.py | 49 ++++++++++++-------------------- 2 files changed, 34 insertions(+), 31 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 30749c187..81dcf4624 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2248,6 +2248,22 @@ def add_branches(self, branches): "number of baskets / basket capacity is not yet supported via the " "cascade approach" ) + if casc._num_baskets > 1: + # add_branches always back-fills a new branch with exactly one + # basket spanning every existing entry (see below), regardless of + # how many baskets the tree's other branches already have. On a + # tree with only one basket that coincidentally matches, but on + # any tree with more, it would silently create the very + # divergent-basket-count state the guard above exists to reject + # -- turning every future extend() on this tree into a permanent + # NotImplementedError, with no warning at add_branches() time. + raise NotImplementedError( + "add_branches for a TTree that already has more than one basket " + "is not yet supported via the cascade approach: the new branch " + "would only ever get one basket, permanently diverging from the " + "other branches' basket counts and blocking any future extend() " + "on this tree" + ) # add new branch dicts to cascade compression = casc._freesegments.fileheader.compression diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index a31e0d3ba..17be27b7e 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -428,26 +428,21 @@ def test_add_branch_does_not_reload_tree_from_disk(tmp_path): assert calls == [] -def test_add_branch_after_multiple_extends(tmp_path): - """add_branches() on a tree that already has more than one basket. - - Regression test: write_updates() stamps every branch's fWriteBasket with - the tree-wide casc._num_baskets, correct for the pre-existing branches - (which really do have that many baskets) but wrong for a brand-new - branch, which add_branches() always writes exactly one basket for. Left - uncorrected, the new branch's fWriteBasket claimed as many baskets as the - rest of the tree even though only basket 0 held real data. That alone - doesn't crash a plain read (this file's new_branch reads back correctly - below), but it makes fWriteBasket agree with the older branches' basket - count despite the real per-branch layout being divergent -- which let a - follow-up extend() sail past the divergent-basket-count guard (since - that guard trusts fWriteBasket) and write another basket for new_branch - indexed as if it were basket 2, when only basket 0 was ever real. Reading - it back then failed with a ValueError about basket/entry counts not - adding up. Every existing add_branches test happened to extend() exactly - once first, so old_num_baskets was coincidentally always 1 and this - never surfaced. With fWriteBasket corrected, the guard now sees the true - divergence and rejects the follow-up extend() cleanly instead. +def test_add_branch_after_multiple_extends_raises(tmp_path): + """add_branches() on a tree that already has more than one basket must refuse up front. + + Regression test: add_branches() always back-fills a new branch with + exactly one basket spanning every existing entry, regardless of how many + baskets the tree's other branches already have. On a tree with only one + basket that coincidentally matches (fWriteBasket needed correcting to 1 + for the new branch -- see test_add_branch_then_extend_same_session, which + covers exactly that case and a follow-up extend() succeeding). But on any + tree with more than one basket, silently proceeding would create the + exact divergent-basket-count state the extend()-side guard exists to + reject -- turning every future extend() on this tree into a permanent + NotImplementedError, with no warning at add_branches() time. Every + existing add_branches test happened to extend() exactly once first, so + this never surfaced. Now add_branches() itself refuses up front instead. """ path = os.path.join(tmp_path, "test.root") with uproot.recreate(path) as f: @@ -456,24 +451,16 @@ def test_add_branch_after_multiple_extends(tmp_path): f["tree"].extend({"x": np.full(100, 2.0, dtype=np.float32)}) with uproot.update(path) as f: - f["tree"].add_branches({"new_branch": np.full(200, 9.0, dtype=np.float32)}) + with pytest.raises(NotImplementedError): + f["tree"].add_branches({"new_branch": np.full(200, 9.0, dtype=np.float32)}) + # the rejected call must not have corrupted the existing branch with uproot.open(path) as f: assert f["tree"].num_entries == 200 - assert np.all(f["tree"]["new_branch"].array() == 9.0) x = f["tree"]["x"].array() assert np.all(x[:100] == 1.0) assert np.all(x[100:] == 2.0) - with uproot.update(path) as f: - with pytest.raises(NotImplementedError): - f["tree"].extend( - { - "x": np.full(50, 3.0, dtype=np.float32), - "new_branch": np.full(50, 8.0, dtype=np.float32), - } - ) - def test_extend_root_written_tree(tmp_path): """extend() on a genuinely ROOT-written tree must not corrupt the file. From 9ea890243c97d217f607bde3b17a2adf1ea9e4bd 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, 26 Aug 2026 17:15:03 +0300 Subject: [PATCH 51/55] fix: recognize non-record awkward Array values as flat-dict extend data, and reject jagged accept_new_fields cleanly --- src/uproot/writing/writable.py | 39 +++++++++++++++++++-- tests/test_1690_ttree_inplace.py | 59 ++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index 81dcf4624..d84ea8bdd 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2445,9 +2445,22 @@ def extend(self, data, *, accept_new_fields=False): 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) + # check if data looks like a flat dict of branch arrays (branch name -> + # array), as opposed to e.g. a single top-level record array. Every + # awkward Array has a `.fields` attribute regardless of its type -- + # it's just empty for a non-record (flat or jagged) array -- so + # `hasattr(v, "fields")` alone can't tell a record apart from a plain + # jagged/flat awkward Array value; checking that `.fields` is + # non-empty is what actually means "this value is itself a record". + # Getting this wrong skipped the new-field detection and + # accept_new_fields auto-add logic below entirely whenever any value + # in data happened to be an awkward Array, jagged or not -- new + # awkward-typed fields fell straight through to the low-level + # extend(), which doesn't know about accept_new_fields and raised a + # confusing "does not correspond to any branch" naming its + # auto-generated counter branch instead. _data_is_flat_dict = isinstance(data, dict) and all( - not hasattr(v, "fields") for v in data.values() + not (hasattr(v, "fields") and v.fields) for v in data.values() ) if isinstance(data, dict) and _data_is_flat_dict: existing_names = _user_branch_names @@ -2484,6 +2497,28 @@ def extend(self, data, *, accept_new_fields=False): "'extend' was given data that do not correspond to any branch: " + repr(next(iter(new_fields))) ) + # add_branches() (used below to back-fill zeros for new fields) + # only creates simple scalar TBranch, with no counter-branch + # support -- it cannot create a new jagged branch. Left + # unchecked, numpy.asarray(v) below raises a confusing, + # awkward-internal "cannot convert to RegularArray" error for + # any jagged new field, instead of saying plainly that this + # isn't supported. + _jagged_new_fields = [ + k + for k, v in new_fields.items() + if isinstance(v, awkward.Array) + and v.ndim > 1 + and not v.layout.purelist_isregular + ] + if _jagged_new_fields: + raise NotImplementedError( + f"accept_new_fields does not yet support jagged (variable-length) " + f"new fields: {_jagged_new_fields}. Call add_branches() with a " + f"scalar zero-filled backfill for this field first (which itself " + f"does not yet support jagged branches either), or restructure the " + f"data to not introduce a jagged field via extend()." + ) zeros = { k: numpy.zeros( self._cascading._num_entries, dtype=numpy.asarray(v).dtype diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 17be27b7e..9b199a211 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -338,6 +338,65 @@ def test_extend_new_fields_error_without_flag(tmp_path): ) +def test_extend_accept_new_fields_flat_awkward_value(tmp_path): + """accept_new_fields must work when the new field's value is an awkward Array. + + Regression test: _data_is_flat_dict checked `not hasattr(v, "fields")` to + decide whether `data` is a plain dict of branch arrays, but every awkward + Array has a `.fields` attribute regardless of type -- it's just empty for + a non-record (flat or jagged) array. That check being always False for + any awkward-Array value skipped the new-field detection and + accept_new_fields auto-add logic entirely, so a new field's data being a + plain (non-jagged) awkward Array fell straight through to the low-level + extend(), which doesn't know about accept_new_fields and raised "does not + correspond to any branch" even though accept_new_fields=True was passed. + """ + ak = pytest.importorskip("awkward") + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(3, dtype=np.float32)}) + + with uproot.update(path) as f: + f["tree"].extend( + {"x": np.ones(2, dtype=np.float32), "y": ak.Array([9.0, 8.0])}, + accept_new_fields=True, + ) + + with uproot.open(path) as f: + assert f["tree"]["x"].array().tolist() == [1.0, 1.0, 1.0, 1.0, 1.0] + assert f["tree"]["y"].array().tolist() == [0.0, 0.0, 0.0, 9.0, 8.0] + + +def test_extend_accept_new_fields_jagged_raises_clearly(tmp_path): + """accept_new_fields with a jagged new field must raise a clear error, not crash confusingly. + + Regression test: add_branches() (used to back-fill zeros for a new + field) only creates simple scalar TBranch, with no counter-branch + support, so it cannot create a jagged branch. Left unchecked, this + reached numpy.asarray(jagged_awkward_array), which raises an + awkward-internal "cannot convert to RegularArray" ValueError -- or, before + the _data_is_flat_dict fix above, an unrelated low-level "'nj', 'j' do + not correspond to any branch" naming an auto-generated counter branch the + user never passed. + """ + ak = pytest.importorskip("awkward") + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(3, dtype=np.float32)}) + + with uproot.update(path) as f: + with pytest.raises(NotImplementedError, match="jagged"): + f["tree"].extend( + { + "x": np.ones(2, dtype=np.float32), + "j": ak.Array([[1.0, 2.0], [3.0]]), + }, + accept_new_fields=True, + ) + + @skip_no_root def test_extend_root_readable(tmp_path): with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: From 6a1e52b80d4db6decff98275475c553f06b46372 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, 27 Aug 2026 10:10:05 +0300 Subject: [PATCH 52/55] fix: opportunistically flatten nested-record extend() input using the default field_name convention when reopened via uproot.update() --- src/uproot/writing/writable.py | 29 ++++++++++++++++++++ tests/test_1690_ttree_inplace.py | 45 ++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index d84ea8bdd..aab497758 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -2463,6 +2463,35 @@ def extend(self, data, *, accept_new_fields=False): not (hasattr(v, "fields") and v.fields) for v in data.values() ) if isinstance(data, dict) and _data_is_flat_dict: + # opportunistically flatten a nested-record-shaped value under a + # key that isn't a recognized branch/record-parent/counter name, + # mirroring mktree's default field_name convention ("outer_inner"). + # A tree reopened via uproot.update() has no persisted + # record-parent metadata to recognize: _load_existing_ttree only + # ever reconstructs "counter"/"normal" branches, because the + # grouping itself was never written to disk -- only the + # already-flattened leaf branches were -- so there is nothing on + # disk to reconstruct it *from* (and field_name is user- + # customizable besides, so branch names alone can't reveal it + # reliably). Rather than guess at the file's original structure, + # expand a dict-valued key only when doing so lines up exactly + # with real existing branches -- this makes the common + # (default-separator) case round-trip the same way extend() + # behaves in the same creation session, without ever silently + # misinterpreting an unrelated key. + _known_top_level = ( + set(_user_branch_names) | _record_names | set(_counter_branch_names) + ) + for _k in list(data.keys()): + if _k in _known_top_level: + continue + _v = data[_k] + if isinstance(_v, Mapping): + _flattened = {f"{_k}_{_sub}": _subv for _sub, _subv in _v.items()} + if _flattened and set(_flattened) <= set(_user_branch_names): + data = {**data, **_flattened} + del data[_k] + existing_names = _user_branch_names # also get record parent names that the user passes as dicts _record_parent_names = { diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 9b199a211..09aacd58e 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -338,6 +338,51 @@ def test_extend_new_fields_error_without_flag(tmp_path): ) +def test_extend_nested_record_same_shape_after_reopen(tmp_path): + """extend() must accept the same nested-dict shape in update mode as at creation. + + Regression test: mktree({"m": {"a": ..., "b": ...}}) creates a "record" + kind entry in _branch_data purely as an in-memory convenience for + extend() to un-nest a dict-shaped value for "m" into its flattened leaf + branches "m_a"/"m_b" -- that grouping is never written to disk, only the + already-flattened leaf branches are. _load_existing_ttree only ever + reconstructs "counter"/"normal" branches, so a tree reopened via + uproot.update() has no "record" entry for "m" at all, and + extend({"m": {...}}) raised "missing: ['m_a', 'm_b']" even though the + identical call worked in the creating session -- forcing the user to + flatten "m" by hand only after reopening the file. Fixed by + opportunistically flattening a dict-shaped value under an unrecognized + key using the default field_name convention ("outer_inner"), but only + when doing so exactly matches real existing branches, so a genuinely + unrelated key or a partial/mismatched record still raises the same clear + error as before. + """ + path = os.path.join(tmp_path, "test.root") + with uproot.recreate(path) as f: + f.mktree("tree", {"m": {"a": np.float64, "b": np.int32}}) + f["tree"].extend( + {"m": {"a": np.array([1.0]), "b": np.array([2], dtype=np.int32)}} + ) + + with uproot.update(path) as f: + # the same nested-dict shape that worked at creation must still work + f["tree"].extend( + {"m": {"a": np.array([3.0]), "b": np.array([4], dtype=np.int32)}} + ) + # the already-flattened form must still work too + f["tree"].extend({"m_a": np.array([5.0]), "m_b": np.array([6], dtype=np.int32)}) + + with uproot.open(path) as f: + assert f["tree"]["m_a"].array().tolist() == [1.0, 3.0, 5.0] + assert f["tree"]["m_b"].array().tolist() == [2, 4, 6] + + with uproot.update(path) as f: + # a genuinely unrelated dict-valued key must still raise clearly, + # not be silently (mis)expanded + with pytest.raises(ValueError, match="missing"): + f["tree"].extend({"nonexistent": {"x": np.array([1.0])}}) + + def test_extend_accept_new_fields_flat_awkward_value(tmp_path): """accept_new_fields must work when the new field's value is an awkward Array. From 51da6129e8a3a24fe66772b5f087c873c4785519 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, 27 Aug 2026 18:11:06 +0300 Subject: [PATCH 53/55] docs: add missing license header to test_1690_ttree_inplace.py --- tests/test_1690_ttree_inplace.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 09aacd58e..2f8d8cd4e 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -1,3 +1,5 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE + import os import shutil From 591e1527a86988cd1636045650cc364e372805ec 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, 27 Aug 2026 18:12:20 +0300 Subject: [PATCH 54/55] test: assert KeyInFileError specifically for add_branches on a nonexistent tree, not just Exception --- tests/test_1690_ttree_inplace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index 2f8d8cd4e..faa540322 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -164,7 +164,7 @@ def test_add_branch_nonexistent_tree(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(Exception): + with pytest.raises(uproot.exceptions.KeyInFileError): f["nonexistent"].add_branches( {"new_branch": np.ones(100, dtype=np.float32)} ) From f2fac6afa660577cef4bd3d741904536070f0ddd 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, 27 Aug 2026 18:13:14 +0300 Subject: [PATCH 55/55] test: assert the specific missing-branches ValueError for extend() with a nonexistent branch, not just Exception --- tests/test_1690_ttree_inplace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py index faa540322..7b14cc8c3 100644 --- a/tests/test_1690_ttree_inplace.py +++ b/tests/test_1690_ttree_inplace.py @@ -299,7 +299,7 @@ def test_extend_nonexistent_branch(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(Exception): + with pytest.raises(ValueError, match="missing"): f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)})