Added some streamer handling

src/uproot/writing/_cascade.py

@@ -713,9 +713,21 @@ def serialize(self, out, branch):
 
         # TODO how to handle this? Make sure to be TBranchElements will be handled too
         # empty TObjArray of TBranches
-        out.append(
-            b"@\x00\x00\x15\x00\x03\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
-        )
+        if len(datum["fBranches"]) == 0:
+            out.append(
+                b"@\x00\x00\x15\x00\x03\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            )
+        else:
+            #     print("serialize branches!!")
+            #     # TObjArray header with fName: ""
+            out.append(b"\x00\x01\x00\x00\x00\x00\x03\x00@\x00\x00")
+            out.append(
+                uproot.models.TObjArray._tobjarray_format1.pack(
+                    len(self._branch_data["fBranches"]),  # TObjArray fSize
+                    0,  # TObjArray fLowerBound
+                )
+            )
+
         subtobjarray_of_leaves_index = len(out)
         out.append(None)
 
@@ -904,7 +916,9 @@ def serialize(self, out, branch):
         # TODO "fBranches, which is a TObjArray of nested TBranch instances (possibly TBranchElement)"
 
         if len(datum["fBaskets"]) >= 1:
-            raise NotImplementedError
+            # print("NotImplementedError, cannot yet write TObjArray of fBaskets")
+            msg = "Cannot yet write baskets"
+            raise NotImplementedError(msg)
 
         out.append(
             b"@\x00\x00\x15\x00\x03\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
@@ -918,11 +932,9 @@ def serialize(self, out, branch):
         # speedbump and fBasketBytes
         out.append(b"\x01")
         out.append(uproot._util.tobytes(datum["fBasketBytes"]))
-
         # speedbump and fBasketEntry
         out.append(b"\x01")
         out.append(uproot._util.tobytes(datum["fBasketEntry"]))
-
         # speedbump and fBasketSeek
         out.append(b"\x01")
         out.append(uproot._util.tobytes(datum["fBasketSeek"]))
@@ -2157,7 +2169,7 @@ def add_branches(
             existing_branches,
             existing_ttree,
         )
-        tree.add_branches(
+        updated_streamers = tree.add_branches(
             sink, directory.file, new_branches
         )  # need new_branches for extend...
         # start = key.seek_location
@@ -2170,7 +2182,7 @@ def add_branches(
         # directory._cascading.write(self._file.sink)
         # directory._file.sink.set_file_length(self._cascading.freesegments.fileheader.end)
         # directory._file.sink.flush()
-        return tree
+        return tree, updated_streamers
 
     def add_rntuple(self, sink, name, title, akform):
         import uproot.writing._cascadentuple

src/uproot/writing/_cascadetree.py

@@ -897,8 +897,7 @@ def write_anew(self, sink):
         num_branches = sum(
             0 if datum["kind"] == "record" else 1 for datum in self._branch_data
         )
-        if self._existing_branches:
-            num_branches += len(self._existing_branches)
+
         # TObjArray header with fName: ""
         out.append(b"\x00\x01\x00\x00\x00\x00\x03\x00@\x00\x00")
         out.append(
@@ -908,16 +907,6 @@ def write_anew(self, sink):
             )
         )
 
-        # Write old branches?
-        if self._existing_branches:
-            old_branches = uproot.writing._cascade.OldBranches(self._existing_branches)
-            for branch in self._existing_branches:
-                # create OldTBranch object
-                # members = uproot.branch.read_members()
-                out, temp = old_branches.serialize(
-                    out, branch
-                )  # should call uproot.models.TBranch._tbranch13_format...pack or something
-                tleaf_reference_numbers.append(temp)  # and don't forget the tleaves
         for datum in self._branch_data:
             if datum["kind"] == "record":
                 continue
@@ -1369,7 +1358,6 @@ def write_np_basket(self, sink, branch_name, compression, array):
         itemsize = array.dtype.itemsize
         for item in array.shape[1:]:
             itemsize *= item
-
         uncompressed_data = uproot._util.tobytes(array)
         compressed_data = uproot.compression.compress(uncompressed_data, compression)
 
@@ -1590,16 +1578,20 @@ def get_tree_key(self):
 
     def add_branches(self, sink, file, new_branches):
         old_key = self.get_tree_key()
-        self.write_with_new_branches(sink, old_key)
-        start = old_key.location
-        stop = start + old_key.num_bytes + old_key.compressed_bytes
-        self._freesegments.release(start, stop)
-        sink.set_file_length(self._freesegments.fileheader.end)
-        sink.flush()
-
+        #     start = old_key.location
+        #     stop = start + old_key.num_bytes + old_key.compressed_bytes
+        #     self._freesegments.release(start, stop)
+        #     sink.set_file_length(self._freesegments.fileheader.end)
+        #     sink.flush()
+        # streamers = [x for x in file._cascading.tlist_of_streamers]
+        streamers = self.write_with_new_branches(sink, old_key)
+        # Reset
+        # old_key = self.get_tree_key()
         self.extend(file, sink, new_branches)
+        return streamers
 
     def write_with_new_branches(self, sink, old_key):
+        models_for_streamers = []
         key_num_bytes = uproot.reading._key_format_big.size + 6
         name_asbytes = self._name.encode(errors="surrogateescape")
         title_asbytes = self._title.encode(errors="surrogateescape")
@@ -1660,6 +1652,7 @@ def write_with_new_branches(self, sink, old_key):
         num_branches = sum(
             0 if datum["kind"] == "record" else 1 for datum in self._branch_data
         )
+        # Include original branches in num_branches
         if self._existing_branches:
             num_branches += len(self._existing_branches)
         # TObjArray header with fName: ""
@@ -1671,16 +1664,13 @@ def write_with_new_branches(self, sink, old_key):
             )
         )
 
-        # Write old branches?
+        # Write old branches
         if self._existing_branches:
             old_branches = uproot.writing._cascade.OldBranches(self._existing_branches)
             for branch in self._existing_branches:
                 # create OldTBranch object
-                # members = uproot.branch.read_members()
-                out, temp = old_branches.serialize(
-                    out, branch
-                )  # should call uproot.models.TBranch._tbranch13_format...pack or something
-                tleaf_reference_numbers.append(temp)  # and don't forget the tleaves
+                out, temp = old_branches.serialize(out, branch)
+                tleaf_reference_numbers.append(temp)
         for datum in self._branch_data:
             if datum["kind"] == "record":
                 continue
@@ -1764,22 +1754,32 @@ def write_with_new_branches(self, sink, old_key):
             out.append(("TLeaf" + letter_upper).encode() + b"\x00")
             if letter_upper == "O":
                 special_struct = uproot.models.TLeaf._tleafO1_format1
+                model = uproot.models.TLeaf.Model_TLeafO_v1.class_rawstreamers
             elif letter_upper == "B":
                 special_struct = uproot.models.TLeaf._tleafb1_format1
+                model = uproot.models.TLeaf.Model_TLeafB_v1
             elif letter_upper == "S":
                 special_struct = uproot.models.TLeaf._tleafs1_format1
+                model = uproot.models.TLeaf.Model_TLeafS_v1
             elif letter_upper == "I":
                 special_struct = uproot.models.TLeaf._tleafi1_format1
+                model = uproot.models.TLeaf.Model_TLeafI_v1
             elif letter_upper == "G":
                 special_struct = uproot.models.TLeaf._tleafl1_format0
             elif letter_upper == "L":
                 special_struct = uproot.models.TLeaf._tleafl1_format0
+                model = uproot.models.TLeaf.Model_TLeafL_v1
             elif letter_upper == "F":
                 special_struct = uproot.models.TLeaf._tleaff1_format1
+                model = uproot.models.TLeaf.Model_TLeafF_v1
             elif letter_upper == "D":
                 special_struct = uproot.models.TLeaf._tleafd1_format1
+                model = uproot.models.TLeaf.Model_TLeafD_v1
             elif letter_upper == "C":
                 special_struct = uproot.models.TLeaf._tleafc1_format1
+                model = uproot.models.TLeaf.Model_TLeafC_v1
+
+            models_for_streamers.append(model)
             fLenType = datum["dtype"].itemsize
             fIsUnsigned = letter != letter_upper
 
@@ -1992,6 +1992,7 @@ def write_with_new_branches(self, sink, old_key):
             replaces=old_key,
             big=True,
         )
+        return models_for_streamers
 
 
 _tbasket_offsets_length = struct.Struct(">I")

src/uproot/writing/writable.py

@@ -1379,7 +1379,7 @@ def add_branches(  # variation of mktree for copying ttree
         names = old_ttree.keys()
         if len(names) == 0:
             raise ValueError(
-                f"""TTree {old_ttree.name} in file {old_ttree.file_path} is empty."""
+                f"""TTree {old_ttree.name} in file {old_ttree.file_path} is empty."""  # TODO does this check need to be here?
             )
         at = -1
         try:  # Will this throw an error? proabably?
@@ -1411,7 +1411,11 @@ def add_branches(  # variation of mktree for copying ttree
                 branch_array = uproot.writing._cascadetree.recarray_to_dict(  # noqa: PLW2901 (overwriting branch_array)
                     branch_array
                 )
-
+            entries = old_ttree.member("fEntries")
+            if len(branch_array) != old_ttree.member("fEntries"):
+                raise ValueError(
+                    f"'add_branches' must fill every branch with the same number of entries; new branches should have {entries} entries, but {branch_name!r} has {len(branch_array)} entries"
+                )
             if isinstance(branch_array, Mapping) and all(
                 isinstance(x, str) for x in branch_array
             ):
@@ -1468,50 +1472,41 @@ def add_branches(  # variation of mktree for copying ttree
                             branch_dtype = numpy.dtype((branch_dtype, branch_shape))
                         metadata[branch_name] = branch_dtype
         file.close()
-        tree = WritableTree(
-            path,
-            directory._file,
-            directory._cascading.add_branches(
-                directory._file.sink,
-                old_ttree.name,
-                old_ttree.title,
-                metadata,
-                counter_name,
-                field_name,
-                initial_basket_capacity,
-                resize_factor,
-                old_ttree,
-                old_ttree.branches,
-                branches,
-                directory,
-            ),
+        update_streamers = []
+        obj, update_streamers = directory._cascading.add_branches(
+            directory._file.sink,
+            old_ttree.name,
+            old_ttree.title,
+            metadata,
+            counter_name,
+            field_name,
+            initial_basket_capacity,
+            resize_factor,
+            old_ttree,
+            old_ttree.branches,
+            branches,
+            directory,
         )
-        # directory._file._new_tree(tree)
-
-        seen = set()
-        streamers = []
-        for model in (
-            uproot.models.TLeaf.Model_TLeafB_v1,
-            uproot.models.TLeaf.Model_TLeafS_v1,
-            uproot.models.TLeaf.Model_TLeafI_v1,
-            uproot.models.TLeaf.Model_TLeafL_v1,
-            uproot.models.TLeaf.Model_TLeafF_v1,
-            uproot.models.TLeaf.Model_TLeafD_v1,
-            uproot.models.TLeaf.Model_TLeafC_v1,
-            uproot.models.TLeaf.Model_TLeafO_v1,
+        tree = WritableTree(path, directory._file, obj)
+        update_streamers.append(
             uproot.models.TBranch.Model_TBranch_v13,
+        )
+        update_streamers.append(
             uproot.models.TTree.Model_TTree_v20,
-        ):
+        )
+        seen = set()
+        streamers = []
+        for model in update_streamers:
             for rawstreamer in model.class_rawstreamers:
                 classname_version = rawstreamer[-2], rawstreamer[-1]
                 if classname_version not in seen:
                     seen.add(classname_version)
                     streamers.append(
                         uproot.writing._cascade.RawStreamerInfo(*rawstreamer)
                     )
-
         directory._file._cascading.streamers.update_streamers(
-            directory._file.sink, streamers
+            directory._file.sink,
+            streamers,
         )
         return tree
 
@@ -1696,7 +1691,6 @@ def update(self, pairs=None, **more_pairs):
         update.
         """
         streamers = []
-
         if pairs is not None:
             if hasattr(pairs, "keys"):
                 all_pairs = itertools.chain(
@@ -1722,13 +1716,9 @@ def update(self, pairs=None, **more_pairs):
                 directory = directory[item]
 
             uproot.writing.identify.add_to_directory(v, name, directory, streamers)
-
         self._file._cascading.streamers.update_streamers(self._file.sink, streamers)
 
 
-# class UpdatableTree:
-
-
 class WritableTree:
     """
     Args:

tests/test_1155_feat_add_copy_ttree.py

@@ -6,8 +6,6 @@
 import pytest
 
 ROOT = pytest.importorskip("ROOT")
-
-# import ROOT
 import numpy as np
 
 import awkward as ak
@@ -36,10 +34,17 @@ def test_vector():
 
 
 def simple_test(tmp_path):
-    import ROOT
-
     data = np.array([1, 2, 3, 4, 5], dtype=np.int64)
-    data1 = np.array([2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.int32)
+    data1 = np.array(
+        [
+            2.0,
+            3.0,
+            4.0,
+            5.0,
+            6.0,
+        ],
+        dtype=np.int32,
+    )
 
     with uproot.recreate(os.path.join(tmp_path, "arrays1.root")) as f:
         f["whatever"] = {"b1": data}
@@ -53,20 +58,67 @@ def simple_test(tmp_path):
     with uproot.open(
         os.path.join(tmp_path, "arrays1.root"), minimal_ttree_metadata=False
     ) as check:
-        # check["tree"].show()
         with uproot.open(
             os.path.join(tmp_path, "arrays2.root"), minimal_ttree_metadata=False
         ) as new:
-            print(new.file.chunk(1358, 2677).raw_data.tobytes(), "\n")
-            print(check.file.chunk(1358, 2677).raw_data.tobytes())
+            inFile = ROOT.TFile.Open(os.path.join(tmp_path, "arrays2.root"), "READ")
+            tree = inFile.Get("whatever;1")
+            print(tree)
+            for x in tree:
+                print(getattr(x, "b1"))
+
+
+def test_subbranches(tmp_path):
+    data = np.array([1, 2, 3, 4, 5], dtype=np.int64)
+    data1 = np.array(
+        [
+            2.0,
+            3.0,
+            4.0,
+            5.0,
+            6.0,
+        ],
+        dtype=np.int32,
+    )
+
+    with uproot.recreate(os.path.join(tmp_path, "arrays2.root")) as f:
+        f["whatever"] = {"b1": data, "b2": data1}
 
+    with uproot.update(os.path.join(tmp_path, "arrays2.root")) as f:
+        f.add_branches("whatever", {"b3": data, "b4": data1})
+
+    with uproot.open(
+        os.path.join(tmp_path, "tree_tester.root"), minimal_ttree_metadata=False
+    ) as check:
+        # check["tree"].show()
+        print(check.keys())
+        with uproot.open(
+            os.path.join(tmp_path, "arrays2.root"), minimal_ttree_metadata=False
+        ) as new:
+            print(new["whatever"].all_members)
             inFile = ROOT.TFile.Open(os.path.join(tmp_path, "arrays2.root"), "READ")
             tree = inFile.Get("whatever;1")
             print(tree)
             for x in tree:
                 print(getattr(x, "b1"))
 
 
+def test_different_fEntries(tmp_path):
+    data = np.array([1, 2, 3, 4, 5], dtype=np.int64)
+    data1 = np.array([2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.int32)
+
+    with uproot.recreate(os.path.join(tmp_path, "arrays2.root")) as f:
+        with pytest.raises(ValueError):
+            f["whatever"] = {"b1": data, "b2": data1}
+            f.add_branches(
+                "whatever",
+                {
+                    "b3": data,
+                    "b4": np.array([2.0, 3.0, 4.0, 5.0, 6.0, 7.0], dtype=np.int32),
+                },
+            )
+
+
 def test_ak_arrays(tmp_path):
 
     data = np.array([1, 2, 3], dtype=np.int64)
@@ -82,13 +134,13 @@ def test_ak_arrays(tmp_path):
     with uproot.recreate(os.path.join(tmp_path, "ak_test.root")) as file:
         file["tree"] = {
             "b1": ak.Array([data, data1, data2]),
-            "b2": ak.Array([data1, data2, data]),
         }
 
     with uproot.update(os.path.join(tmp_path, "ak_test.root")) as write:
         write.add_branches(
             "tree",
             {
+                "b2": ak.Array([data1, data2, data]),
                 "b3": ak.Array([data2, data, data1]),
             },
         )
@@ -99,29 +151,173 @@ def test_ak_arrays(tmp_path):
         with uproot.open(
             os.path.join(tmp_path, "ak_test.root"), minimal_ttree_metadata=False
         ) as new:
-
-            import ROOT
-
             inFile = ROOT.TFile.Open(os.path.join(tmp_path, "ak_test.root"), "READ")
             tree = inFile.Get("tree")
-            print(tree.GetBranch("b1"))
             for x in tree:
-
                 print(getattr(x, "b1"))
-            print(tree.GetBranch("b2"))
-            for x in tree:
-                print(getattr(x, "b2"))
+            print(tree.Scan())
+            # ak.Array()
+            # for x in tree:
+            #     print(getattr(x, "b2").GetArray())
+
+
+def test_streamers_same_dtypes(tmp_path):
+    from ROOT import TTree
+    from array import array
+
+    N = 4
+    data = array("f", N * [0.0])
+    data1 = array("f", [2.0, 3.0, 4.0, 5.0])
+
+    inFile = root.TFile(
+        "/Users/zobil/Desktop/directory/root_streamers_F.root", "RECREATE"
+    )
+    tree = root.TTree("tree1", "tree")
+    import numpy as np
+
+    # Basic type branch (float) - use array of length 1
+    # n = array('f', [ 1.5 ])
+    # tree.Branch('b1', n, 'b1/F')
+
+    # Array branch - use array of length N
+    N = 4
+    # a = array('d', N*[ 0. ])
+    # tree.Branch('b1', a, 'b1[' + str(N) + ']/D')
+
+    # # Array branch - use NumPy array of length N
+    npa = np.zeros(4, dtype=np.float32)
+    tree.Branch("b1", npa, "b1/F")
+    for i in range(4):
+        npa[0] = i**0
+        tree.Fill()
+    inFile.Write()
+    inFile.Close()
+
+    inFile = root.TFile.Open(os.path.join(tmp_path, "root_streamers_F.root"), "OPEN")
+    tree = inFile.Get("tree1")
+    tree.Scan()
+    data = np.array([5.0, 6.0, 7.0, 8.0], dtype=np.float32)
+
+    with uproot.update(os.path.join(tmp_path, "root_streamers_F.root")) as file:
+        file.add_branches("tree1", {"b2": data})
+
+    with uproot.open(
+        os.path.join(tmp_path, "root_streamers_F.root"), minimal_ttree_metadata=False
+    ) as file:
+        inFile = ROOT.TFile.Open(
+            os.path.join(tmp_path, "root_streamers_F.root"), "READ"
+        )
+        tree = inFile.Get("tree1;1")
+        indx = 0
+        for x in tree:
+            assert getattr(x, "b1") == file["tree1"]["b1"].array()[indx]
+            assert getattr(x, "b2") == file["tree1"]["b2"].array()[indx]
+            indx += 1
+
+        tree.Scan()
+        check = [
+            "TBranch",
+            "TAttLine",
+            "TCollection",
+            "TLeafF",
+            "listOfRules",
+            "TString",
+            "TObjArray",
+            "TAttFill",
+            "TBranchRef",
+            "TList",
+            "ROOT::TIOFeatures",
+            "TSeqCollection",
+            "TAttMarker",
+            "TTree",
+            "TNamed",
+            "TObject",
+            "TAttLine",
+            "TLeaf",
+            "TRefTable",
+        ]
+        for i in set(file.file.streamers):
+            assert i in check
+
+
+def test_streamers_diff_dtypes(tmp_path):
+
+    inFile = ROOT.TFile(
+        "/Users/zobil/Desktop/directory/root_diff_dtypes.root", "RECREATE"
+    )
+    tree = ROOT.TTree("tree1", "tree")
+
+    # Basic type branch (float) - use array of length 1
+    # n = array('f', [ 1.5 ])
+    # tree.Branch('b1', n, 'b1/F')
+
+    # Array branch - use array of length N
+    N = 4
+    # a = array('d', N*[ 0. ])
+    # tree.Branch('b1', a, 'b1[' + str(N) + ']/D')
+
+    # # Array branch - use NumPy array of length N
+    npa = np.zeros(4, dtype=float)
+    tree.Branch("b1", npa, "b1F")
+    for i in range(4):
+        npa[0] = i**0
+        tree.Fill()
+    inFile.Write()
+    inFile.Close()
+
+    inFile = ROOT.TFile.Open(os.path.join(tmp_path, "root_diff_dtypes.root"), "OPEN")
+    tree = inFile.Get("tree1")
+    tree.Scan()
+    data = np.array([5, 6, 7, 8], dtype=np.int64)
+    with uproot.update(os.path.join(tmp_path, "root_diff_dtypes.root")) as file:
+        file.add_branches("tree1", {"b2": data})
+
+    with uproot.open(
+        os.path.join(tmp_path, "root_diff_dtypes.root"), minimal_ttree_metadata=False
+    ) as file:
+        file["tree1"]["b2"].member("fLeaves")[0].all_members
+
+        inFile = ROOT.TFile.Open(
+            os.path.join(tmp_path, "root_diff_dtypes.root"), "READ"
+        )
+        tree = inFile.Get("tree1;1")
+        indx = 0
+        for x in tree:
+            assert getattr(x, "b1") == file["tree1"]["b1"].array()[indx]
+            assert getattr(x, "b2") == file["tree1"]["b2"].array()[indx]
+            indx += 1
+
+        # tree.Scan()
+        check = [
+            "TBranch",
+            "TAttLine",
+            "TCollection",
+            "TLeafF",
+            "listOfRules",
+            "TString",
+            "TObjArray",
+            "TAttFill",
+            "TBranchRef",
+            "TList",
+            "ROOT::TIOFeatures",
+            "TSeqCollection",
+            "TAttMarker",
+            "TTree",
+            "TNamed",
+            "TObject",
+            "TAttLine",
+            "TLeaf",
+            "TRefTable",
+            "TLeafL",
+        ]
+        for i in set(file.file.streamers):
+            assert i in check
 
 
 def HZZ_test(tmp_path):
     with uproot.open(
         data_path("uproot-HZZ.root"), minimal_ttree_metadata=False
     ) as test:
-
-        # print(test["events"]["NMuon"].typename)
-        # print(test["events"])
-        # print(test['events'].all_members)
-
         with uproot.update(os.path.join(tmp_path, "uproot-HZZ.root copy")) as new:
             # data = np.arange(0, 2420, 1)
             data = []
@@ -134,30 +330,13 @@ def HZZ_test(tmp_path):
             os.path.join(tmp_path, "uproot-HZZ.root copy"),
             minimal_ttree_metadata=False,
         ) as check:
-            print(check.keys(cycle=False))
-            print(check["events"]["data"].array())
-            print(check["events"].arrays())
-            print(test["events"].arrays())
-
             for key in test["events"].keys():
                 assert key in test["events"].keys()
                 assert ak.all(
                     check["events"][key].array() == test["events"][key].array()
                 )
 
-            # print(check.file.chunk.start, check['events'].chunk.stop)
-            # print(check.file.chunk.get(1000, 2000, check['events'].cursor, context=None).tobytes())
-
-            import ROOT
-
             inFile = ROOT.TFile.Open(
                 os.path.join(tmp_path, "uproot-HZZ.root copy"), "READ"
             )
             tree = inFile.Get("events")
-
-        # print(check["events"]["Photon_Px"].member("fLeaves")[0].member("fLeafCount"))
-
-
-simple_test("/Users/zobil/Desktop/directory")
-# HZZ_test("/Users/zobil/Desktop/directory")
-# test_ak_arrays("/Users/zobil/Desktop/directory")