feat: Add in-place TTree branch addition and row extension

src/uproot/writing/writable.py

@@ -1014,9 +1014,8 @@ 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"
-                )
+                # 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)
@@ -1054,6 +1053,195 @@ 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())
+                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(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)
+                        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
+                    )
+                    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",
+                ]
+            }
+        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 = 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
+
+        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:
@@ -1883,10 +2071,121 @@ def num_baskets(self) -> int:
         """
         return self._cascading.num_baskets
 
-    def extend(self, data):
+    def add_branches(self, branches):
+        """
+        Args:
+            branches (dict of str -> array): Names and data of new 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)})
+        """
+        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
+        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
+
+        for branch_name, branch_data in branches.items():
+            arr = numpy.asarray(branch_data)
+            if len(arr) != num_entries:
+                raise ValueError(
+                    f"branch {branch_name!r} has {len(arr)} entries but TTree has "
+                    f"{num_entries} entries; all new branches must match the tree length"
+                )
+            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:
+            raise NotImplementedError(
+                "add_branches for files with TBranchElement branches is not yet "
+                "supported via the cascade approach"
+            )
+
+        # 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"]
+            )
+            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()
+
+        # 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
+
+        # update self._cascading so subsequent extend uses correct metadata
+        writable_tree = uproot.writing.writable.WritableTree(
+            self._path, self._file, casc
+        )
+        self._file._trees[casc._key.seek_location] = writable_tree
+        self._cascading = casc
+
+    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`.
@@ -1910,6 +2209,66 @@ 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 <https://numpy.org/doc/stable/reference/generated/numpy.ndarray.nbytes.html>`__ 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) <https://github.com/scikit-hep/uproot5/pull/428#issuecomment-908703486>`__.
         """
+        if self._cascading is None:
+            raise RuntimeError(
+                "_cascading is None — this should not happen; please report this bug"
+            )
+        # validate branches
+        # 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.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(
+                    f"'extend' must fill every branch with the same number of entries; missing: {missing}"
+                )
+            if new_fields:
+                if not accept_new_fields:
+                    raise ValueError(
+                        "'extend' was given data that do not correspond to any branch: "
+                        + repr(next(iter(new_fields)))
+                    )
+                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(