diff --git a/uproot/interp/auto.py b/uproot/interp/auto.py index 34787823..3d4b3c3f 100644 --- a/uproot/interp/auto.py +++ b/uproot/interp/auto.py @@ -323,7 +323,7 @@ def transform(node, tofloat=True): if branch._isTClonesArray: return asgenobj(SimpleArray(STLVector(asdtype(fromdtype, fromdtype))), branch._context, 6) else: - return asjagged(ascontent, skipbytes=10) + return asjagged(ascontent, skipbytes=10, sizeat=6) except _NotNumerical: if branch._vecstreamer is not None: @@ -348,9 +348,9 @@ def transform(node, tofloat=True): return asgenobj(STLVector(streamerClass), branch._context, 6) else: if streamerClass._methods is None: - return asjagged(astable(asdtype(recarray)), skipbytes=10) + return asjagged(astable(asdtype(recarray)), skipbytes=10, sizeat=6) else: - return asjagged(asobj(astable(asdtype(recarray)), streamerClass._methods), skipbytes=10) + return asjagged(asobj(astable(asdtype(recarray)), streamerClass._methods), skipbytes=10, sizeat=6) if hasattr(branch._streamer, "_fTypeName"): m = re.match(b"bitset<([1-9][0-9]*)>", branch._streamer._fTypeName) @@ -358,27 +358,27 @@ def transform(node, tofloat=True): return asjagged(asstlbitset(int(m.group(1))), skipbytes=6) if getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype(awkward.numpy.bool_), skipbytes=10) + return asjagged(asdtype(awkward.numpy.bool_), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("i1"), skipbytes=10) + return asjagged(asdtype("i1"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("u1"), skipbytes=10) + return asjagged(asdtype("u1"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("i2"), skipbytes=10) + return asjagged(asdtype("i2"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("u2"), skipbytes=10) + return asjagged(asdtype("u2"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("i4"), skipbytes=10) + return asjagged(asdtype("i4"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("u4"), skipbytes=10) + return asjagged(asdtype("u4"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("i8"), skipbytes=10) + return asjagged(asdtype("i8"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("u8"), skipbytes=10) + return asjagged(asdtype("u8"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("f4"), skipbytes=10) + return asjagged(asdtype("f4"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector" or getattr(branch._streamer, "_fTypeName", None) == b"vector": - return asjagged(asdtype("f8"), skipbytes=10) + return asjagged(asdtype("f8"), skipbytes=10, sizeat=6) elif getattr(branch._streamer, "_fTypeName", None) == b"vector": return asgenobj(STLVector(STLString(awkward)), branch._context, 6) else: @@ -445,27 +445,27 @@ def transform(node, tofloat=True): return asstring(skipbytes=1) if branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype(awkward.numpy.bool_), skipbytes=10) + return asjagged(asdtype(awkward.numpy.bool_), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("i1"), skipbytes=10) + return asjagged(asdtype("i1"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("u1"), skipbytes=10) + return asjagged(asdtype("u1"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("i2"), skipbytes=10) + return asjagged(asdtype("i2"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("u2"), skipbytes=10) + return asjagged(asdtype("u2"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("i4"), skipbytes=10) + return asjagged(asdtype("i4"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("u4"), skipbytes=10) + return asjagged(asdtype("u4"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("i8"), skipbytes=10) + return asjagged(asdtype("i8"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("u8"), skipbytes=10) + return asjagged(asdtype("u8"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("f4"), skipbytes=10) + return asjagged(asdtype("f4"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector" or branch._fClassName == b"vector": - return asjagged(asdtype("f8"), skipbytes=10) + return asjagged(asdtype("f8"), skipbytes=10, sizeat=6) elif branch._fClassName == b"vector": return asgenobj(STLVector(STLString(awkward)), branch._context, 6) diff --git a/uproot/interp/jagged.py b/uproot/interp/jagged.py index 97495827..5a4431bb 100644 --- a/uproot/interp/jagged.py +++ b/uproot/interp/jagged.py @@ -31,12 +31,13 @@ class asjagged(uproot.interp.interp.Interpretation): # makes __doc__ attribute mutable before Python 3.3 __metaclass__ = type.__new__(type, "type", (uproot.interp.interp.Interpretation.__metaclass__,), {}) - def __init__(self, content, skipbytes=0): + def __init__(self, content, skipbytes=0, sizeat=None): self.content = content self.skipbytes = skipbytes + self.sizeat = sizeat def __repr__(self): - return "asjagged({0}{1})".format(repr(self.content), "" if self.skipbytes == 0 else ", {0}".format(self.skipbytes)) + return "asjagged({0}{1}{2})".format(repr(self.content), "" if self.skipbytes == 0 else ", {0}".format(self.skipbytes), "" if self.sizeat == 0 else ", {0}".format(self.sizeat)) def to(self, todtype=None, todims=None, skipbytes=None): if skipbytes is None: @@ -75,6 +76,16 @@ def fromroot(self, data, byteoffsets, local_entrystart, local_entrystop, keylen) return self.awkward.JaggedArray(starts, stops, content) else: + if self.sizeat is not None: + sizeat_bytestarts = byteoffsets[local_entrystart : local_entrystop] + self.sizeat + sizeat_good_bytestarts = sizeat_bytestarts[sizeat_bytestarts + 4 <= len(data)] + sizeat_mask = self.awkward.numpy.zeros(len(data), dtype=self.awkward.numpy.int8) + sizeat_mask[sizeat_good_bytestarts + 0] = 1 + sizeat_mask[sizeat_good_bytestarts + 1] = 1 + sizeat_mask[sizeat_good_bytestarts + 2] = 1 + sizeat_mask[sizeat_good_bytestarts + 3] = 1 + sizeat = data[sizeat_mask.view(self.awkward.numpy.bool_)].view(">i4") + bytestarts = byteoffsets[local_entrystart : local_entrystop ] + self.skipbytes bytestops = byteoffsets[local_entrystart + 1 : local_entrystop + 1] @@ -106,7 +117,14 @@ def fromroot(self, data, byteoffsets, local_entrystart, local_entrystop, keylen) offsets[0] = 0 self.awkward.numpy.cumsum(counts, out=offsets[1:]) - return self.awkward.JaggedArray(offsets[:-1], offsets[1:], content) + starts = offsets[:-1] + if self.sizeat is not None: + stops = starts + sizeat + if not (stops == offsets[1:]).all(): + return self.awkward.JaggedArray(starts, stops, content).compact() + + stops = offsets[1:] + return self.awkward.JaggedArray(starts, stops, content) def destination(self, numitems, numentries): content = self.content.destination(numitems, numentries) diff --git a/uproot/version.py b/uproot/version.py index 10bec281..6bd77250 100644 --- a/uproot/version.py +++ b/uproot/version.py @@ -6,7 +6,7 @@ import re -__version__ = "3.11.3" +__version__ = "3.11.4" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))