Merge pull request #326 from scikit-hep/issue-324

Implemented TRef as an object that can live in TTrees.
scikit-hep · Aug 30, 2019 · 3a12730 · 3a12730
2 parents a1a87c8 + 918e838
commit 3a12730
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 11 deletions.
diff --git a/uproot/interp/auto.py b/uproot/interp/auto.py
@@ -118,19 +118,16 @@ def _obj_or_genobj(streamerClass, branch, isjagged, cntvers=False, tobject=True,
             return asgenobj(streamerClass, context, 0)
 
     else:
-        if streamerClass._methods is None:
-            return asdtype(recarray)
+        if isjagged:
+            if streamerClass._methods is None:
+                return asjagged(astable(asdtype(recarray)))
+            else:
+                return asjagged(asobj(astable(asdtype(recarray)), streamerClass._methods))
         else:
-            if isjagged:
-                if streamerClass._methods is None:
-                    return asjagged(astable(asdtype(recarray)))
-                else:
-                    return asjagged(asobj(astable(asdtype(recarray)), streamerClass._methods))
+            if streamerClass._methods is None:
+                return asdtype(recarray)
             else:
-                if streamerClass._methods is None:
-                    return asdtype(recarray)
-                else:
-                    return asobj(astable(asdtype(recarray)), streamerClass._methods)
+                return asobj(astable(asdtype(recarray)), streamerClass._methods)
 
 def interpret(branch, awkwardlib=None, swapbytes=True, cntvers=False, tobject=True, speedbump=True):
     import uproot.tree

diff --git a/uproot/rootio.py b/uproot/rootio.py
@@ -1324,6 +1324,32 @@ def _readinto(cls, self, source, cursor, context, parent):
         TList._readinto(self, source, cursor, context, parent)
         return self
 
+class TRef(ROOTStreamedObject):
+    _format1 = struct.Struct(">xxIxxxxxx")
+
+    def __init__(self, id):
+        self.id = id
+
+    @classmethod
+    def _readinto(cls, self, source, cursor, context, parent):
+        self.id = cursor.field(source, self._format1)
+        return self
+
+    def __repr__(self):
+        return "<TRef {0}>".format(self.id)
+
+    @classmethod
+    def _recarray(cls):
+        out = []
+        out.append(("pidf", ">u2"))
+        out.append(("id", ">u4"))
+        out.append((" other", "S6"))
+        return out
+
+TRef._methods = TRef
+TRef._arraymethods = None
+TRef._fromrow = lambda row: TRef(row["id"])
+
 class TArray(list, ROOTStreamedObject):
     @classmethod
     def _readinto(cls, self, source, cursor, context, parent):
@@ -1419,6 +1445,7 @@ def __repr__(self):
                    "TObjString":     TObjString,
                    "TList":          TList,
                    "THashList":      THashList,
+                   "TRef":           TRef,
                    "TArray":         TArray,
                    "TArrayC":        TArrayC,
                    "TArrayS":        TArrayS,