pypy / pypy / module / cppyy / capi /

Full commit
import py, os, sys

from pypy.interpreter.error import OperationError
from pypy.interpreter.gateway import interp2app, unwrap_spec
from pypy.interpreter.typedef import TypeDef
from pypy.interpreter.baseobjspace import Wrappable

from pypy.translator.tool.cbuild import ExternalCompilationInfo
from pypy.rpython.lltypesystem import rffi
from pypy.rlib import libffi, rdynload

from pypy.module.itertools import interp_itertools

__all__ = ['identify', 'eci', 'c_load_dictionary']

pkgpath = py.path.local(__file__).dirpath().join(os.pardir)
srcpath = pkgpath.join("src")
incpath = pkgpath.join("include")

if os.environ.get("ROOTSYS"):
    import commands
    (stat, incdir) = commands.getstatusoutput("root-config --incdir")
    if stat != 0:        # presumably Reflex-only
        rootincpath = [os.path.join(os.environ["ROOTSYS"], "include")]
        rootlibpath = [os.path.join(os.environ["ROOTSYS"], "lib64"), os.path.join(os.environ["ROOTSYS"], "lib")]
        rootincpath = [incdir]
        rootlibpath = commands.getoutput("root-config --libdir").split()
    rootincpath = []
    rootlibpath = []

def identify():
    return 'CINT'

ts_reflect = False
ts_call    = False
ts_memory  = 'auto'
ts_helper  = 'auto'

# force loading in global mode of core libraries, rather than linking with
# them as PyPy uses various version of dlopen in various places; note that
# this isn't going to fly on Windows (note that locking them in objects and
# calling dlclose in __del__ seems to come too late, so this'll do for now)
with rffi.scoped_str2charp('') as ll_libname:
    _cintdll = rdynload.dlopen(ll_libname, rdynload.RTLD_GLOBAL | rdynload.RTLD_NOW)
with rffi.scoped_str2charp('') as ll_libname:
    _coredll = rdynload.dlopen(ll_libname, rdynload.RTLD_GLOBAL | rdynload.RTLD_NOW)

eci = ExternalCompilationInfo(
    include_dirs=[incpath] + rootincpath,
    link_extra=["-lCore", "-lCint"],

_c_load_dictionary = rffi.llexternal(
    [rffi.CCHARP], rdynload.DLLHANDLE,

def c_load_dictionary(name):
    result = _c_load_dictionary(name)
    if not result:
        err = rdynload.dlerror()
        raise rdynload.DLOpenError(err)
    return libffi.CDLL(name)       # should return handle to already open file

# CINT-specific pythonizations ===============================================

### TTree --------------------------------------------------------------------
_ttree_Branch = rffi.llexternal(
    [rffi.VOIDP, rffi.CCHARP, rffi.CCHARP, rffi.VOIDP, rffi.INT, rffi.INT], rffi.LONG,

def ttree_Branch(space, w_self, args_w):
    """Pythonized version of TTree::Branch(): takes proxy objects and by-passes
    the CINT-manual layer."""

    from pypy.module.cppyy import interp_cppyy
    tree_class = interp_cppyy.scope_byname(space, "TTree")

    # sigs to modify (and by-pass CINT):
    #  1. (const char*, const char*, T**,               Int_t=32000, Int_t=99)
    #  2. (const char*, T**,                            Int_t=32000, Int_t=99)
    argc = len(args_w)

    # basic error handling of wrong arguments is best left to the original call,
    # so that error messages etc. remain consistent in appearance: the following
    # block may raise TypeError or IndexError to break out anytime

        if argc < 2 or 5 < argc:
            raise TypeError("wrong number of arguments")

        tree = space.interp_w(interp_cppyy.W_CPPInstance, w_self, can_be_None=True)
        if (tree is None) or (tree.cppclass != tree_class):
            raise TypeError("not a TTree")

        # first argument must always always be cont char*
        branchname = space.str_w(args_w[0])

        # if args_w[1] is a classname, then case 1, else case 2
            classname = space.str_w(args_w[1])
            addr_idx  = 2
            w_address = args_w[addr_idx]
        except OperationError:
            addr_idx  = 1
            w_address = args_w[addr_idx]

        bufsize, splitlevel = 32000, 99
        if addr_idx+1 < argc: bufsize = space.c_int_w(args_w[addr_idx+1])
        if addr_idx+2 < argc: splitlevel = space.c_int_w(args_w[addr_idx+2])

        # now retrieve the W_CPPInstance and build other stub arguments
        space =    # holds the class cache in State
        cppinstance = space.interp_w(interp_cppyy.W_CPPInstance, w_address)
        address = rffi.cast(rffi.VOIDP, cppinstance.get_rawobject())
        klassname = cppinstance.cppclass.full_name()
        vtree = rffi.cast(rffi.VOIDP, tree.get_rawobject())

        # call the helper stub to by-pass CINT
        vbranch = _ttree_Branch(vtree, branchname, klassname, address, bufsize, splitlevel)
        branch_class = interp_cppyy.scope_byname(space, "TBranch")
        w_branch = interp_cppyy.wrap_cppobject(
            space, space.w_None, branch_class, vbranch, isref=False, python_owns=False)
        return w_branch
    except (OperationError, TypeError, IndexError), e:

    # 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")

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 =            # 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):

        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

        self.getentry = tree.cppclass.get_overload("GetEntry").functions[0]
        self.current  = 0
        self.maxentry = space.int_w(space.call_method(w_tree, "GetEntriesFast"))

        space = =          # holds the class cache in State
        space.call_method(w_tree, "SetBranchStatus", space.wrap("*"), space.wrap(0))

    def iter_w(self):

    def next_w(self):
        if self.current == self.maxentry:
            raise OperationError(,
        # TODO: check bytes read?, [])
        self.current += 1 
        return self.w_tree

W_TTreeIter.typedef = TypeDef(
    __iter__ = interp2app(W_TTreeIter.iter_w),
    next = interp2app(W_TTreeIter.next_w),

def ttree_iter(space, w_self):
    """Allow iteration over TTree's. Also initializes branch data members and
    sets addresses, if needed."""
    w_treeiter = W_TTreeIter(space, w_self)
    return w_treeiter

# setup pythonizations for later use at run-time
_pythonizations = {}
def register_pythonizations(space):

    ### TTree
    _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):

    if name == 'TFile':
        space.setattr(w_pycppclass, space.wrap("__getattr__"),
                      space.getattr(w_pycppclass, space.wrap("Get")))

    elif name == 'TTree':
        space.setattr(w_pycppclass, space.wrap("_unpythonized_Branch"),
                      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__"),
                      space.getattr(w_pycppclass, space.wrap("GetNoElements")))