1. jaeyeun
  2. pypy

Commits

wlav  committed 4d1ca3a

optimized I/O for CINT backend

  • Participants
  • Parent commits a10072d
  • Branches reflex-support

Comments (0)

Files changed (2)

File pypy/module/cppyy/capi/cint_capi.py

View file
     # return control back to the original, unpythonized overload
     return tree_class.get_overload("Branch").call(w_self, args_w)
 
+def activate_branch(space, w_branch):
+    w_branches = space.call_method(w_branch, "GetListOfBranches")
+    for i in range(space.int_w(space.call_method(w_branches, "GetEntriesFast"))):
+        w_b = space.call_method(w_branches, "At", space.wrap(i))
+        activate_branch(space, w_b)
+    space.call_method(w_branch, "SetStatus", space.wrap(1))
+    space.call_method(w_branch, "ResetReadEntry")
+
+@unwrap_spec(args_w='args_w')
+def ttree_getattr(space, w_self, args_w):
+    """Specialized __getattr__ for TTree's that allows switching on/off the
+    reading of individual branchs."""
+
+    from pypy.module.cppyy import interp_cppyy
+    tree = space.interp_w(interp_cppyy.W_CPPInstance, w_self)
+
+    # setup branch as a data member and enable it for reading
+    space = tree.space            # holds the class cache in State
+    w_branch = space.call_method(w_self, "GetBranch", args_w[0])
+    w_klassname = space.call_method(w_branch, "GetClassName")
+    klass = interp_cppyy.scope_byname(space, space.str_w(w_klassname))
+    w_obj = klass.construct()
+    #space.call_method(w_branch, "SetStatus", space.wrap(1)) 
+    activate_branch(space, w_branch)
+    space.call_method(w_branch, "SetObject", w_obj)
+    space.call_method(w_branch, "GetEntry", space.wrap(0))
+    space.setattr(w_self, args_w[0], w_obj)
+    return w_obj
+
 class W_TTreeIter(Wrappable):
     def __init__(self, space, w_tree):
-        self.current = 0
+
+        from pypy.module.cppyy import interp_cppyy
+        tree = space.interp_w(interp_cppyy.W_CPPInstance, w_tree)
+        self.tree = tree.get_cppthis(tree.cppclass)
         self.w_tree = w_tree
-        from pypy.module.cppyy import interp_cppyy
-        tree = space.interp_w(interp_cppyy.W_CPPInstance, self.w_tree)
-        self.tree = tree.get_cppthis(tree.cppclass)
+
         self.getentry = tree.cppclass.get_overload("GetEntry").functions[0]
+        self.current  = 0
+        self.maxentry = space.int_w(space.call_method(w_tree, "GetEntriesFast"))
 
-        # setup data members if this is the first iteration time
-        try:
-            space.getattr(w_tree, space.wrap("_pythonized"))
-        except OperationError:
-            self.space = space = tree.space       # holds the class cache in State
-            w_branches = space.call_method(w_tree, "GetListOfBranches")
-            for i in range(space.int_w(space.call_method(w_branches, "GetEntriesFast"))):
-                w_branch = space.call_method(w_branches, "At", space.wrap(i))
-                w_name = space.call_method(w_branch, "GetName")
-                w_klassname = space.call_method(w_branch, "GetClassName")
-                klass = interp_cppyy.scope_byname(space, space.str_w(w_klassname))
-                w_obj = klass.construct()
-                space.call_method(w_branch, "SetObject", w_obj)
-                # cache the object and define this tree pythonized
-                space.setattr(w_tree, w_name, w_obj)
-                space.setattr(w_tree, space.wrap("_pythonized"), space.w_True)
+        space = self.space = tree.space          # holds the class cache in State
+        space.call_method(w_tree, "SetBranchStatus", space.wrap("*"), space.wrap(0))
 
     def iter_w(self):
         return self.space.wrap(self)
 
     def next_w(self):
-        w_bytes_read = self.getentry.call(self.tree, [self.space.wrap(self.current)])
-        if not self.space.is_true(w_bytes_read):
+        if self.current == self.maxentry:
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
+        # TODO: check bytes read?
+        self.getentry.call(self.tree, [self.space.wrap(self.current)])
         self.current += 1 
         return self.w_tree
 
     "NOT_RPYTHON"
 
     ### TTree
-    _pythonizations['ttree_Branch'] = space.wrap(interp2app(ttree_Branch))
-    _pythonizations['ttree_iter']   = space.wrap(interp2app(ttree_iter))
+    _pythonizations['ttree_Branch']  = space.wrap(interp2app(ttree_Branch))
+    _pythonizations['ttree_iter']    = space.wrap(interp2app(ttree_iter))
+    _pythonizations['ttree_getattr'] = space.wrap(interp2app(ttree_getattr))
 
 # callback coming in when app-level bound classes have been created
 def pythonize(space, name, w_pycppclass):
                       space.getattr(w_pycppclass, space.wrap("Branch")))
         space.setattr(w_pycppclass, space.wrap("Branch"), _pythonizations["ttree_Branch"])
         space.setattr(w_pycppclass, space.wrap("__iter__"), _pythonizations["ttree_iter"])
+        space.setattr(w_pycppclass, space.wrap("__getattr__"), _pythonizations["ttree_getattr"])
 
     elif name[0:8] == "TVectorT":    # TVectorT<> template
         space.setattr(w_pycppclass, space.wrap("__len__"),

File pypy/module/cppyy/test/test_cint.py

View file
             import cppyy
             return cppyy.load_reflection_info(%r)""" % (iotypes_dct,))
 
-    def test01_write_stdvector( self ):
+    def test01_write_stdvector(self):
         """Test writing of a single branched TTree with an std::vector<double>"""
 
         from cppyy import gbl               # bootstraps, only needed for tests
 
         i = 0
         for event in mytree:
+            assert len(event.mydata) == self.M
             for entry in event.mydata:
                 assert i == int(entry)
                 i += 1
         f = TFile(self.fname)
         mytree = f.Get(self.tname)
 
+        j = 1
         for event in mytree:
             i = 0
+            assert len(event.data.get_floats()) == j*self.M
             for entry in event.data.get_floats():
                 assert i == int(entry)
                 i += 1
 
+            k = 1
+            assert len(event.data.get_tuples()) == j
             for mytuple in event.data.get_tuples():
                 i = 0
+                assert len(mytuple) == k*self.M
                 for entry in mytuple:
                     assert i == int(entry)
                     i += 1
+                k += 1
+            j += 1
+        assert j-1 == self.N
         #
         f.Close()
+
+    def test05_branch_activation(self):
+        """Test of automatic branch activation"""
+
+        from cppyy import gbl               # bootstraps, only needed for tests
+        from cppyy.gbl import TFile, TTree
+        from cppyy.gbl.std import vector
+
+        L = 5
+
+        # writing
+        f = TFile(self.fname, "RECREATE")
+        mytree = TTree(self.tname, self.title)
+        mytree._python_owns = False
+
+        for i in range(L):
+            v = vector("double")()
+            mytree.Branch("mydata_%d"%i, v.__class__.__name__, v)
+            mytree.__dict__["v_%d"%i] = v
+
+        for i in range(self.N):
+            for k in range(L):
+                v = mytree.__dict__["v_%d"%k]
+                for j in range(self.M):
+                    mytree.__dict__["v_%d"%k].push_back(i*self.M+j*L+k)
+            mytree.Fill()
+            for k in range(L):
+                v = mytree.__dict__["v_%d"%k]
+                v.clear()
+        f.Write()
+        f.Close()
+
+        del mytree, f
+        import gc
+        gc.collect()
+
+        # reading
+        f = TFile(self.fname)
+        mytree = f.Get(self.tname)
+
+        # force (initial) disabling of all branches
+        mytree.SetBranchStatus("*",0);
+
+        i = 0
+        for event in mytree:
+            for k in range(L):
+                j = 0
+                data = getattr(mytree, "mydata_%d"%k)
+                assert len(data) == self.M
+                for entry in data:
+                    assert entry == i*self.M+j*L+k
+                    j += 1
+                assert j == self.M
+            i += 1
+        assert i == self.N
+