Skip to content

Commit ed6f4aa

Browse files
zbilodeazbilodea
committed
Copies, details to work on though and haven't filled branches yet
1 parent f137305 commit ed6f4aa

File tree

3 files changed

+453
-2
lines changed

3 files changed

+453
-2
lines changed

‎src/uproot/writing/_cascade.py

+340
Original file line numberDiff line numberDiff line change
@@ -580,6 +580,344 @@ def deserialize(cls, raw_bytes, location, num_bytes, num_slices, in_path):
580580
return out
581581

582582

583+
class OldBranch(CascadeLeaf): # Branch or branches?
584+
"""
585+
A :doc:`uproot.writing._cascade.CascadeLeaf` for copying an old TBranch to a new TTree. ?
586+
"""
587+
588+
class_version = 1
589+
590+
def __init__(self, branch_data):
591+
self._branch_data = branch_data
592+
593+
@property
594+
def allocation(self):
595+
if self._allocation is None:
596+
self._allocation = self.num_bytes
597+
return self._allocation
598+
599+
@allocation.setter
600+
def allocation(self, value):
601+
if self._allocation != value:
602+
self._file_dirty = True
603+
self._allocation = value
604+
605+
@property
606+
def num_bytes(self):
607+
total = 0
608+
for _, stop in self._slices:
609+
if stop - 1 >= uproot.const.kStartBigFile:
610+
total += _free_format_big.size
611+
else:
612+
total += _free_format_small.size
613+
614+
if self._end is None:
615+
if total + _free_format_small.size >= uproot.const.kStartBigFile:
616+
total += _free_format_big.size
617+
else:
618+
total += _free_format_small.size
619+
elif self._end >= uproot.const.kStartBigFile:
620+
total += _free_format_big.size
621+
else:
622+
total += _free_format_small.size
623+
624+
return total
625+
626+
def serialize(self, out):
627+
# superclass TNamed (Model_TNamed(uproot.model.Model))
628+
# superclass TAttFill
629+
key_num_bytes = uproot.reading._key_format_big.size + 6
630+
name_asbytes = self._branch_data.tree.name.encode(errors="surrogateescape")
631+
title_asbytes = self._branch_data.tree.title.encode(errors="surrogateescape")
632+
key_num_bytes += (1 if len(name_asbytes) < 255 else 5) + len(name_asbytes)
633+
key_num_bytes += (1 if len(title_asbytes) < 255 else 5) + len(title_asbytes)
634+
635+
any_tbranch_index = len(out)
636+
out.append(None)
637+
# if isinstance(self._branch_data, uproot.models.TBranchElement):
638+
639+
out.append(b"TBranch\x00")
640+
641+
tbranch_index = len(out)
642+
out.append(None)
643+
644+
tbranch_tobject = uproot.models.TObject.Model_TObject.empty() # ?
645+
# tbranch_tnamed = self._branch_data['TNamed'].serialize() # ?
646+
tbranch_tnamed = uproot.models.TNamed.Model_TNamed.empty()
647+
tbranch_tnamed._bases.append(tbranch_tobject)
648+
tbranch_tnamed._members["fTitle"] = self._branch_data.title
649+
tbranch_tnamed._serialize(
650+
out, True, self._branch_data.name, numpy.uint32(0x00400000)
651+
)
652+
653+
# TAttFill v2, fFillColor: 0, fFillStyle: 1001
654+
# make model TAttFill v2 with fFillColor and fFillStyle
655+
tattfill = uproot.models.TAtt.Model_TAttFill_v2.empty()
656+
# tattfill._deeply_writable = True # ?
657+
tattfill._members["fFillColor"] = self._branch_data.member("fFillColor")
658+
tattfill._members["fFillStyle"] = self._branch_data.member("fFillStyle")
659+
660+
out.append(tattfill.serialize(out))
661+
662+
self._branch_data.members["metadata_start"] = (6 + 6 + 8 + 6) + sum(
663+
len(x) for x in out if x is not None
664+
)
665+
666+
# Lie about the compression level so that ROOT checks and does the right thing.
667+
# https://github.com/root-project/root/blob/87a998d48803bc207288d90038e60ff148827664/tree/tree/src/TBasket.cxx#L560-L578
668+
# Without this, when small buffers are left uncompressed, ROOT complains about them not being compressed.
669+
# (I don't know where the "no, really, this is uncompressed" bit is.)
670+
671+
# Have to actually make something for if there's a TBranchElement!!
672+
673+
out.append(
674+
uproot.models.TBranch._tbranch13_format1.pack(
675+
self._branch_data.member("fCompress"),
676+
self._branch_data.member("fBasketSize"),
677+
self._branch_data.member("fEntryOffsetLen"),
678+
self._branch_data.member("fWriteBasket"), # fWriteBasket
679+
self._branch_data.member("fEntryNumber"), # fEntryNumber
680+
)
681+
)
682+
683+
# fIOFeatures (TIOFeatures)
684+
out.append(b"@\x00\x00\x07\x00\x00\x1a\xa1/\x10\x00")
685+
# 0 to bytestring??
686+
out.append(
687+
uproot.models.TBranch._tbranch13_format2.pack(
688+
self._branch_data.member("fOffset"),
689+
self._branch_data.member("fMaxBaskets"), # fMaxBaskets
690+
self._branch_data.member("fSplitLevel"),
691+
self._branch_data.member("fEntries"), # fEntries
692+
self._branch_data.member("fFirstEntry"),
693+
self._branch_data.member("fTotBytes"),
694+
self._branch_data.member("fZipBytes"),
695+
)
696+
)
697+
698+
# empty TObjArray of TBranches
699+
out.append(
700+
self._branch_data.member("fBranches").serialize(
701+
out,
702+
)
703+
)
704+
705+
subtobjarray_of_leaves_index = len(out)
706+
out.append(None)
707+
708+
# TObjArray header with fName: "", fSize: 1, fLowerBound: 0
709+
out.append(
710+
b"\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00"
711+
)
712+
_dtype_to_char = {
713+
numpy.dtype("bool"): "O",
714+
numpy.dtype(">i1"): "B",
715+
numpy.dtype(">u1"): "b",
716+
numpy.dtype(">i2"): "S",
717+
numpy.dtype(">u2"): "s",
718+
numpy.dtype(">i4"): "I",
719+
numpy.dtype(">u4"): "i",
720+
numpy.dtype(">i8"): "L",
721+
numpy.dtype(">u8"): "l",
722+
numpy.dtype(">f4"): "F",
723+
numpy.dtype(">f8"): "D",
724+
numpy.dtype(">U"): "C",
725+
}
726+
727+
absolute_location = key_num_bytes + sum(len(x) for x in out if x is not None)
728+
absolute_location += 8 + 6 * (sum(1 if x is None else 0 for x in out) - 1)
729+
tleaf_reference_number = absolute_location + 2
730+
731+
subany_tleaf_index = len(out)
732+
out.append(None)
733+
for leaf in self._branch_data.member("fLeaves"):
734+
# Make and serialize each leaf??
735+
letter_upper = _dtype_to_char[numpy.dtype(">i8")]
736+
out.append(("TLeaf" + letter_upper).encode() + b"\x00")
737+
if letter_upper == "O":
738+
special_struct = uproot.models.TLeaf._tleafO1_format1
739+
elif letter_upper == "B":
740+
special_struct = uproot.models.TLeaf._tleafb1_format1
741+
elif letter_upper == "S":
742+
special_struct = uproot.models.TLeaf._tleafs1_format1
743+
elif letter_upper == "I":
744+
special_struct = uproot.models.TLeaf._tleafi1_format1
745+
elif letter_upper == "G":
746+
special_struct = uproot.models.TLeaf._tleafl1_format0
747+
elif letter_upper == "L":
748+
special_struct = uproot.models.TLeaf._tleafl1_format0
749+
elif letter_upper == "F":
750+
special_struct = uproot.models.TLeaf._tleaff1_format1
751+
elif letter_upper == "D":
752+
special_struct = uproot.models.TLeaf._tleafd1_format1
753+
elif letter_upper == "C":
754+
special_struct = uproot.models.TLeaf._tleafc1_format1
755+
# single TLeaf
756+
757+
leaf_name = self._branch_data.member("fName").encode(
758+
errors="surrogateescape"
759+
)
760+
leaf_title = (
761+
self._branch_data.member("fLeaves")[0]
762+
.member("fTitle")
763+
.encode(errors="surrogateescape")
764+
)
765+
leaf_name_length = (1 if len(leaf_name) < 255 else 5) + len(leaf_name)
766+
leaf_title_length = (1 if len(leaf_title) < 255 else 5) + len(leaf_title)
767+
768+
leaf_header = numpy.array(
769+
[
770+
64,
771+
0,
772+
0,
773+
76,
774+
0,
775+
1,
776+
64,
777+
0,
778+
0,
779+
54,
780+
0,
781+
2,
782+
64,
783+
0,
784+
0,
785+
30,
786+
0,
787+
1,
788+
0,
789+
1,
790+
0,
791+
0,
792+
0,
793+
0,
794+
3,
795+
0,
796+
0,
797+
0,
798+
],
799+
numpy.uint8,
800+
)
801+
tmp = leaf_header[0:4].view(">u4")
802+
tmp[:] = (
803+
numpy.uint32(
804+
42 + leaf_name_length + leaf_title_length + special_struct.size
805+
)
806+
| uproot.const.kByteCountMask
807+
)
808+
tmp = leaf_header[6:10].view(">u4")
809+
tmp[:] = (
810+
numpy.uint32(36 + leaf_name_length + leaf_title_length)
811+
| uproot.const.kByteCountMask
812+
)
813+
tmp = leaf_header[12:16].view(">u4")
814+
tmp[:] = (
815+
numpy.uint32(12 + leaf_name_length + leaf_title_length)
816+
| uproot.const.kByteCountMask
817+
)
818+
819+
out.append(uproot._util.tobytes(leaf_header))
820+
if len(leaf_name) < 255:
821+
out.append(
822+
struct.pack(">B%ds" % len(leaf_name), len(leaf_name), leaf_name)
823+
)
824+
else:
825+
out.append(
826+
struct.pack(
827+
">BI%ds" % len(leaf_name), 255, len(leaf_name), leaf_name
828+
)
829+
)
830+
if len(leaf_title) < 255:
831+
out.append(
832+
struct.pack(">B%ds" % len(leaf_title), len(leaf_title), leaf_title)
833+
)
834+
else:
835+
out.append(
836+
struct.pack(
837+
">BI%ds" % len(leaf_title), 255, len(leaf_title), leaf_title
838+
)
839+
)
840+
841+
out.append(
842+
uproot.models.TLeaf._tleaf2_format0.pack(
843+
leaf.member("fLen"),
844+
leaf.member("fLenType"),
845+
leaf.member("fOffset"), # fOffset
846+
leaf.member("fIsRange"), # fIsRange
847+
leaf.member("fIsUnsigned"),
848+
)
849+
)
850+
out.append(
851+
uproot.serialization.serialize_object_any(
852+
leaf.member("fLeafCount") # fLeafCount
853+
)
854+
)
855+
856+
# specialized TLeaf* members (fMinimum, fMaximum)
857+
# datum["tleaf_special_struct"] = special_struct
858+
out.append(
859+
special_struct.pack(
860+
int(leaf.member("fMinimum")), int(leaf.member("fMaximum"))
861+
)
862+
)
863+
864+
out[subany_tleaf_index] = (
865+
uproot.serialization._serialize_object_any_format1.pack(
866+
numpy.uint32(sum(len(x) for x in out[subany_tleaf_index + 1 :]) + 4)
867+
| uproot.const.kByteCountMask,
868+
uproot.const.kNewClassTag,
869+
)
870+
)
871+
872+
out[subtobjarray_of_leaves_index] = uproot.serialization.numbytes_version(
873+
sum(len(x) for x in out[subtobjarray_of_leaves_index + 1 :]),
874+
3, # TObjArray
875+
)
876+
877+
# empty TObjArray of fBaskets (embedded)
878+
# TODO "fBranches, which is a TObjArray of nested TBranch instances (possibly TBranchElement)"
879+
880+
out.append(
881+
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"
882+
)
883+
884+
# out.append(self._branch_data.member(""))
885+
886+
# assert sum(1 if x is None else 0 for x in out) == 4
887+
self._branch_data.members["basket_metadata_start"] = (6 + 6 + 8 + 6) + sum(
888+
len(x) for x in out if x is not None
889+
) # ?
890+
891+
# speedbump and fBasketBytes
892+
out.append(b"\x01")
893+
out.append(uproot._util.tobytes(self._branch_data.member("fBasketBytes")))
894+
895+
# speedbump and fBasketEntry
896+
out.append(b"\x01")
897+
out.append(uproot._util.tobytes(self._branch_data.member("fBasketEntry")))
898+
899+
# speedbump and fBasketSeek
900+
out.append(b"\x01")
901+
out.append(uproot._util.tobytes(self._branch_data.member("fBasketSeek")))
902+
903+
self._branch_data.member("fFileName").serialize() # name = None?
904+
905+
out[tbranch_index] = uproot.serialization.numbytes_version(
906+
sum(len(x) for x in out[tbranch_index + 1 :]), 13 # TBranch
907+
)
908+
909+
out[any_tbranch_index] = (
910+
uproot.serialization._serialize_object_any_format1.pack(
911+
numpy.uint32(sum(len(x) for x in out[any_tbranch_index + 1 :]) + 4)
912+
| uproot.const.kByteCountMask,
913+
uproot.const.kNewClassTag,
914+
)
915+
)
916+
return out, tleaf_reference_number
917+
918+
# def write(self, )
919+
920+
583921
class FreeSegments(CascadeNode):
584922
"""
585923
A :doc:`uproot.writing._cascade.CascadeNode` for writing a ROOT FreeSegments record.
@@ -1710,6 +2048,7 @@ def add_tree(
17102048
field_name,
17112049
initial_basket_capacity,
17122050
resize_factor,
2051+
existing_branches=None,
17132052
):
17142053
import uproot.writing._cascadetree
17152054

@@ -1723,6 +2062,7 @@ def add_tree(
17232062
field_name,
17242063
initial_basket_capacity,
17252064
resize_factor,
2065+
existing_branches,
17262066
)
17272067
tree.write_anew(sink)
17282068
return tree

0 commit comments

Comments
 (0)