Merge pull request #245 from scikit-hep/issue-243

Add STLMap as a genobj type

tests/test_issues.py

@@ -211,3 +211,8 @@ def test_issue240(self):
         else:
             t = uproot.open("root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root")["Events"]
             assert (abs(t.array("nMuon")) < 50).all()
+
+    def test_issue243(self):
+        t = uproot.open("tests/samples/issue243.root")["triggerList"]
+        for x in t.array("triggerMap", entrystop=100):
+            assert all(y == 1.0 for y in x.values())

uproot/__init__.py

@@ -160,6 +160,7 @@
 from uproot.interp.objects import asstring
 from uproot.interp.objects import SimpleArray
 from uproot.interp.objects import STLVector
+from uproot.interp.objects import STLMap
 from uproot.interp.objects import STLString
 asdebug = asjagged(asdtype("u1"))
 
@@ -174,4 +175,4 @@
 # don't expose uproot.uproot; it's ugly
 del uproot
 
-__all__ = ["open", "xrootd", "http", "iterate", "numentries", "lazyarray", "lazyarrays", "daskarray", "daskarrays", "daskframe", "create", "recreate", "update", "MemmapSource", "FileSource", "XRootDSource", "HTTPSource", "interpret", "asdtype", "asarray", "asdouble32", "asstlbitset", "asjagged", "astable", "asobj", "asgenobj", "asstring", "asdebug", "SimpleArray", "STLVector", "STLString", "__version__"]
+__all__ = ["open", "xrootd", "http", "iterate", "numentries", "lazyarray", "lazyarrays", "daskarray", "daskarrays", "daskframe", "create", "recreate", "update", "MemmapSource", "FileSource", "XRootDSource", "HTTPSource", "interpret", "asdtype", "asarray", "asdouble32", "asstlbitset", "asjagged", "astable", "asobj", "asgenobj", "asstring", "asdebug", "SimpleArray", "STLVector", "STLMap", "STLString", "__version__"]

uproot/interp/auto.py

@@ -46,6 +46,7 @@
 from uproot.interp.objects import asstring
 from uproot.interp.objects import SimpleArray
 from uproot.interp.objects import STLVector
+from uproot.interp.objects import STLMap
 from uproot.interp.objects import STLString
 
 class _NotNumerical(Exception): pass
@@ -379,6 +380,31 @@ def transform(node, tofloat=True):
                 elif getattr(branch._streamer, "_fTypeName", None) == b"vector<string>":
                     return asgenobj(STLVector(STLString(awkward)), branch._context, 6)
 
+                if getattr(branch._streamer, "_fTypeName", None) == b"map<string,bool>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype(awkward.numpy.bool_)), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,char>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("i1")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,unsigned char>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("u1")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,short>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("i2")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,unsigned short>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("u2")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,int>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("i4")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,unsigned int>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("u4")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,long>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("i8")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,unsigned long>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("u8")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,float>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("f4")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,double>":
+                    return asgenobj(STLMap(STLString(awkward), asdtype("f8")), branch._context, 6)
+                elif getattr(branch._streamer, "_fTypeName", None) == b"map<string,string>":
+                    return asgenobj(STLMap(STLString(awkward), STLString(awkward)), branch._context, 6)
+
                 if getattr(branch._streamer, "_fTypeName", None) == b"vector<vector<bool> >":
                     return asgenobj(STLVector(STLVector(asdtype(awkward.numpy.bool_))), branch._context, 6)
                 elif getattr(branch._streamer, "_fTypeName", None) == b"vector<vector<char> >":

uproot/interp/objects.py

@@ -78,6 +78,43 @@ def read(self, source, cursor, context, parent):
                 out[i] = self.cls.read(source, cursor, context, parent)
             return out
 
+class STLMap(object):
+    def __init__(self, keycls, valcls):
+        self.keycls = keycls
+        self.valcls = valcls
+
+    def __repr__(self):
+        key = self.keycls.__name__ if isinstance(self.keycls, type) else repr(self.keycls)
+        val = self.valcls.__name__ if isinstance(self.valcls, type) else repr(self.valcls)
+        return "STLMap({0}, {1})".format(key, val)
+
+    _format1 = struct.Struct(">i")
+
+    def read(self, source, cursor, context, parent):
+        numitems = cursor.field(source, self._format1)
+
+        out = {}
+        for i in range(numitems):
+            if isinstance(self.keycls, uproot.interp.numerical.asdtype):
+                key = cursor.array(source, 1, self.keycls.fromdtype)
+                if key.dtype != self.keycls.todtype:
+                    key = key.astype(self.keycls.todtype)
+                key = key[0]
+            else:
+                key = self.keycls.read(source, cursor, context, parent)
+
+            if isinstance(self.valcls, uproot.interp.numerical.asdtype):
+                val = cursor.array(source, 1, self.valcls.fromdtype)
+                if val.dtype != self.valcls.todtype:
+                    val = val.astype(self.valcls.todtype)
+                val = val[0]
+            else:
+                val = self.valcls.read(source, cursor, context, parent)
+
+            out[key] = val
+
+        return out
+
 class STLString(object):
     def __init__(self, awkward=None):
         if awkward is None:
@@ -88,9 +125,12 @@ def __repr__(self):
         return "STLString()"
 
     _format1 = struct.Struct("B")
+    _format2 = struct.Struct(">i")
 
     def read(self, source, cursor, context, parent):
         numitems = cursor.field(source, self._format1)
+        if numitems == 255:
+            numitems = cursor.field(source, self._format2)
         return cursor.array(source, numitems, self.awkward.ObjectArray.CHARTYPE).tostring()
 
 class astable(uproot.interp.interp.Interpretation):

uproot/version.py

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