fixes #26: finally added std::bitset<N>

Author: jpivarski

uproot/interp/auto.py

@@ -36,6 +36,7 @@
 import uproot.const
 from uproot.interp.numerical import asdtype
 from uproot.interp.numerical import asarray
+from uproot.interp.numerical import asstlbitset
 from uproot.interp.jagged import asjagged
 from uproot.interp.jagged import asstlvector
 from uproot.interp.jagged import asstlvectorvector
@@ -204,6 +205,11 @@ def interpret(branch, classes=None, swapbytes=True):
                     except _NotNumerical:
                         pass
 
+                if hasattr(branch._streamer, "fTypeName"):
+                    m = re.match(b"bitset<([1-9][0-9]*)>", branch._streamer.fTypeName)
+                    if m is not None:
+                        return asstlbitset(int(m.group(1)))
+
                 if getattr(branch._streamer, "fTypeName", None) == b"vector<bool>":
                     return asstlvector(asdtype("bool"))
                 elif getattr(branch._streamer, "fTypeName", None) == b"vector<char>":
@@ -252,6 +258,10 @@ def interpret(branch, classes=None, swapbytes=True):
                 elif getattr(branch._streamer, "fTypeName", None) == b"vector<vector<double> >":
                     return asstlvectorvector(">f8")
 
+                m = re.match(b"bitset<([1-9][0-9]*)>", branch.fClassName)
+                if m is not None:
+                    return asstlbitset(int(m.group(1)))
+
                 if branch.fClassName == b"vector<bool>":
                     return asstlvector(asdtype("bool"))
                 elif branch.fClassName == b"vector<char>":

uproot/interp/numerical.py

@@ -39,6 +39,12 @@
 else:
     string_types = (str, bytes)
 
+def _dimsprod(dims):
+    out = 1
+    for x in dims:
+        out *= x
+    return out
+
 class asdtype(Interpretation):
     # makes __doc__ attribute mutable before Python 3.3
     __metaclass__ = type.__new__(type, "type", (Interpretation.__metaclass__,), {})
@@ -100,7 +106,8 @@ def empty(self):
         return numpy.empty((0,) + self.todims, dtype=self.todtype)
 
     def compatible(self, other):
-        return isinstance(other, (asdtype, asarray)) and self.todtype == other.todtype and self.todims == other.todims
+        return (isinstance(other, (asdtype, asarray)) and self.todtype == other.todtype and self.todims == other.todims) or \
+               (isinstance(other, asstlbitset) and self.todtype == other.dtype and self.todims == (other.numbytes,))
 
     def numitems(self, numbytes, numentries):
         return numbytes // self.fromdtype.itemsize
@@ -140,6 +147,8 @@ def clip(self, destination, itemstart, itemstop, entrystart, entrystop):
         assert itemstart % product == 0
         assert itemstop % product == 0
         return destination[itemstart // product : itemstop // product]
+        # FIXME: isn't the above equivalent to the following?
+        #     return destination[entrystart:entrystop]
 
     def finalize(self, destination):
         return destination
@@ -195,8 +204,46 @@ def finalize(self, destination):
         array, stop = destination
         return array[:stop]
 
-def _dimsprod(dims):
-    out = 1
-    for x in dims:
-        out *= x
-    return out
+class asstlbitset(Interpretation):
+    # makes __doc__ attribute mutable before Python 3.3
+    __metaclass__ = type.__new__(type, "type", (Interpretation.__metaclass__,), {})
+
+    dtype = numpy.dtype(numpy.bool_)
+
+    def __init__(self, numbytes):
+        self.numbytes = numbytes
+
+    def __repr__(self):
+        return self.identifier
+
+    @property
+    def identifier(self):
+        return "asstlbitset({0})".format(self.numbytes)
+
+    def empty(self):
+        return numpy.empty((0,), dtype=self.dtype)
+
+    def compatible(self, other):
+        return (isinstance(other, asstlbitset) and self.numbytes == other.numbytes) or \
+               (isinstance(other, (asdtype, asarray)) and self.dtype == other.todtype and (self.numbytes,) == other.todims)
+
+    def numitems(self, numbytes, numentries):
+        return max(0, numbytes // self.dtype.itemsize - 10*numentries)
+
+    def source_numitems(self, source):
+        return _dimsprod(source.shape)
+
+    def fromroot(self, data, offsets, local_entrystart, local_entrystop):
+        return data.view(self.dtype).reshape((-1, self.numbytes + 10))[local_entrystart:local_entrystop,10:]
+
+    def destination(self, numitems, numentries):
+        return numpy.empty((numitems // self.numbytes, self.numbytes), dtype=self.dtype)
+
+    def fill(self, source, destination, itemstart, itemstop, entrystart, entrystop):
+        destination[entrystart:entrystop] = source
+
+    def clip(self, destination, itemstart, itemstop, entrystart, entrystop):
+        return destination[entrystart:entrystop]
+
+    def finalize(self, destination):
+        return destination

uproot/tree.py

@@ -552,10 +552,19 @@ def _normalize_branches(self, arg):
     def _normalize_entrystartstop(self, entrystart, entrystop):
         if entrystart is None:
             entrystart = 0
+        elif entrystart < 0:
+            entrystart += self.numentries
+        entrystart = min(self.numentries, max(0, entrystart))
+
         if entrystop is None:
             entrystop = self.numentries
+        elif entrystop < 0:
+            entrystop += self.numentries
+        entrystop = min(self.numentries, max(0, entrystop))
+
         if entrystop < entrystart:
             raise IndexError("entrystop must be greater than or equal to entrystart")
+
         return entrystart, entrystop
 
     def __len__(self):

uproot/version.py

@@ -30,7 +30,7 @@
 
 import re
 
-__version__ = "2.6.7"
+__version__ = "2.6.8"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))