Distinguish between 1-byte and 6-byte header in AsArray. (#84)

scikit-hep · Aug 26, 2020 · ccf790f · ccf790f
1 parent d40ded6
commit ccf790f
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 48 deletions.
diff --git a/uproot4/__init__.py b/uproot4/__init__.py
@@ -62,7 +62,7 @@
 import uproot4.interpretation
 import uproot4.interpretation.library
 from uproot4.interpretation.numerical import AsDtype
-from uproot4.interpretation.numerical import AsArray
+from uproot4.interpretation.numerical import AsDtypeInPlace
 from uproot4.interpretation.numerical import AsDouble32
 from uproot4.interpretation.numerical import AsFloat16
 from uproot4.interpretation.numerical import AsSTLBits

diff --git a/uproot4/behaviors/TBranch.py b/uproot4/behaviors/TBranch.py
@@ -1385,7 +1385,11 @@ def streamer(self):
                                         break
                                 break
 
-                    if self.parent.member("fClassName") == "TClonesArray":
+                    if (
+                        self.parent.member("fClassName") == "TClonesArray"
+                        or self.parent.member("fClonesName", none_if_missing=True)
+                        == fParentName
+                    ):
                         self._streamer_isTClonesArray = True
 
             elif fClassName is not None and fClassName != "":

diff --git a/uproot4/containers.py b/uproot4/containers.py
@@ -328,10 +328,15 @@ def __eq__(self, other):
 
 
 class AsArray(AsContainer):
-    def __init__(self, header, values):
+    def __init__(self, header, speedbump, values):
         self._header = header
+        self._speedbump = speedbump
         self._values = values
 
+    @property
+    def speedbump(self):
+        return self._speedbump
+
     @property
     def values(self):
         return self._values
@@ -341,11 +346,13 @@ def __repr__(self):
             values = self._values.__name__
         else:
             values = repr(self._values)
-        return "AsArray({0}, {1})".format(self.header, values)
+        return "AsArray({0}, {1}, {2})".format(self.header, self.speedbump, values)
 
     @property
     def cache_key(self):
-        return "AsArray({0},{1})".format(self.header, _content_cache_key(self._values))
+        return "AsArray({0},{1},{2})".format(
+            self.header, self.speedbump, _content_cache_key(self._values)
+        )
 
     @property
     def typename(self):
@@ -359,24 +366,66 @@ def awkward_form(self, file, index_format="i64", header=False, tobject_header=Tr
             uproot4._util.awkward_form(
                 self._values, file, index_format, header, tobject_header
             ),
-            parameters={"uproot": {"as": "array", "header": self._header}},
+            parameters={
+                "uproot": {
+                    "as": "array",
+                    "header": self._header,
+                    "speedbump": self._speedbump,
+                }
+            },
         )
 
     def read(self, chunk, cursor, context, file, selffile, parent, header=True):
         if self._header and header:
-            cursor.skip(1)
+            start_cursor = cursor.copy()
+            num_bytes, instance_version = uproot4.deserialization.numbytes_version(
+                chunk, cursor, context
+            )
 
-        if isinstance(self._values, numpy.dtype):
-            remainder = chunk.remainder(cursor.index, cursor, context)
-            return remainder.view(self._values)
+            if isinstance(self._values, numpy.dtype):
+                remainder = chunk.get(
+                    cursor.index, cursor.index + num_bytes, cursor, context
+                )
+                return remainder.view(self._values)
+
+            else:
+                out = []
+                while cursor.displacement(start_cursor) < num_bytes:
+                    out.append(
+                        self._values.read(
+                            chunk, cursor, context, file, selffile, parent
+                        )
+                    )
+
+                if self._header and header:
+                    uproot4.deserialization.numbytes_check(
+                        chunk,
+                        start_cursor,
+                        cursor,
+                        num_bytes,
+                        self.typename,
+                        context,
+                        file.file_path,
+                    )
+                return numpy.array(out, dtype=numpy.dtype(numpy.object))
 
         else:
-            out = []
-            while cursor.index < chunk.stop:
-                out.append(
-                    self._values.read(chunk, cursor, context, file, selffile, parent)
-                )
-            return numpy.array(out, dtype=numpy.dtype(numpy.object))
+            if self._speedbump:
+                cursor.skip(1)
+
+            if isinstance(self._values, numpy.dtype):
+                remainder = chunk.remainder(cursor.index, cursor, context)
+                return remainder.view(self._values)
+
+            else:
+                out = []
+                while cursor.index < chunk.stop:
+                    out.append(
+                        self._values.read(
+                            chunk, cursor, context, file, selffile, parent
+                        )
+                    )
+                return numpy.array(out, dtype=numpy.dtype(numpy.object))
 
 
 class AsDynamic(AsContainer):