Skip to content

Commit

Permalink
[PyROOT] Implement TTree SetBranchAddress pythonization in Python
Browse files Browse the repository at this point in the history
The pythonization of `TTree::SetBranchAddress` was implemented in C++,
hacking into CPyCppy by using implementation details like data member
caches (this call: `((CPPInstance *)address)GetDatamemberCache()`). Not
too surprising that it apparently breaks with the upcoming Python 3.13.

It's better to implement the pythonizations in Python and also manage
the lifetime of the necessary data in Python. This is done in this
commit.

The pythonization is extensively tested in `ttree_setbranchaddress.py`.

Possible closes #15799.
  • Loading branch information
guitargeek committed Jun 13, 2024
1 parent 28f7ba6 commit 5b0c854
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 74 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@
*/
'''

from libROOTPythonizations import GetBranchAttr, SetBranchAddressPyz, BranchPyz
from libROOTPythonizations import GetBranchAttr, BranchPyz
from . import pythonization

# TTree iterator
Expand All @@ -145,15 +145,66 @@ def _TTree__iter__(self):
if bytes_read == -1:
raise RuntimeError("TTree I/O error")

def _SetBranchAddress(self, *args):
# Modify the behaviour if args is (const char*, void*)
res = SetBranchAddressPyz(self, *args)
def _pythonize_branch_addr(branch, addr_orig):
"""Helper for the SetBranchAddress pythonization, extracting the relevant
address from a Python object if possible.
"""
import cppyy
import ctypes

if res is None:
# Fall back to the original implementation for the rest of overloads
res = self._OriginalSetBranchAddress(*args)
# Pythonization for cppyy proxies (of type CPPInstance)
if isinstance(addr_orig, cppyy._backend.CPPInstance):

return res
is_leaf_list = branch.IsA() is cppyy.gbl.TBranch.Class()

if is_leaf_list:
# If the branch is a leaf list, SetBranchAddress expects the
# address of the object that has the corresponding data members.
return ctypes.c_void_p(cppyy.addressof(instance=addr_orig, byref=False))

# Otherwise, SetBranchAddress is expecting a pointer to the address of
# the object, and the pointer needs to stay alive. Therefore, we create
# a container for the pointer and cache it in the original cppyy proxy.
addr_view = cppyy.gbl.array["std::intptr_t", 1]([cppyy.addressof(instance=addr_orig, byref=False)])

if not hasattr(addr_orig, "_set_branch_cached_pointers"):
addr_orig._set_branch_cached_pointers = []
addr_orig._set_branch_cached_pointers.append(addr_view)

# Finally, we have to return the address of the container
return ctypes.c_void_p(cppyy.addressof(instance=addr_view, byref=False))

# For NumPy arrays
if hasattr(addr_orig, "__array_interface__"):
return ctypes.c_void_p(addr_orig.__array_interface__["data"][0])

# For the builtin array library
if hasattr(addr_orig, "buffer_info"):
return ctypes.c_void_p(addr_orig.buffer_info()[0])

# We don't know how to pythonize the address parameter. return the
# original value one.
return addr_orig

def _SetBranchAddress(self, bname, addr, *args, **kwargs):
"""
Pythonization for TTree::SetBranchAddress.
Modify the behaviour of SetBranchAddress so that proxy references can be passed
as arguments from the Python side, more precisely in cases where the C++
implementation of the method expects the address of a pointer.
For example:
```
v = ROOT.std.vector('int')()
t.SetBranchAddress("my_vector_branch", v)
```
"""

branch = self.GetBranch(bname)
addr = _pythonize_branch_addr(branch, addr)

return self._OriginalSetBranchAddress(bname, addr, *args, **kwargs)

def _Branch(self, *args):
# Modify the behaviour if args is one of:
Expand Down
2 changes: 0 additions & 2 deletions bindings/pyroot/pythonizations/src/PyROOTModule.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,6 @@ static PyMethodDef gPyROOTMethods[] = {
(char *)"Cast the void* returned by TClass::DynamicCast to the right type"},
{(char *)"AddTObjectEqNePyz", (PyCFunction)PyROOT::AddTObjectEqNePyz, METH_VARARGS,
(char *)"Add equality and inequality comparison operators to TObject"},
{(char *)"SetBranchAddressPyz", (PyCFunction)PyROOT::SetBranchAddressPyz, METH_VARARGS,
(char *)"Fully enable the use of TTree::SetBranchAddress from Python"},
{(char *)"BranchPyz", (PyCFunction)PyROOT::BranchPyz, METH_VARARGS,
(char *)"Fully enable the use of TTree::Branch from Python"},
{(char *)"AddPrettyPrintingPyz", (PyCFunction)PyROOT::AddPrettyPrintingPyz, METH_VARARGS,
Expand Down
1 change: 0 additions & 1 deletion bindings/pyroot/pythonizations/src/PyROOTPythonize.h
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ PyObject *AddPrettyPrintingPyz(PyObject *self, PyObject *args);

PyObject *GetBranchAttr(PyObject *self, PyObject *args);
PyObject *BranchPyz(PyObject *self, PyObject *args);
PyObject *SetBranchAddressPyz(PyObject *self, PyObject *args);

PyObject *AddTClassDynamicCastPyz(PyObject *self, PyObject *args);

Expand Down
63 changes: 0 additions & 63 deletions bindings/pyroot/pythonizations/src/TTreePyz.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -201,69 +201,6 @@ PyObject *PyROOT::GetBranchAttr(PyObject * /*self*/, PyObject *args)
return 0;
}

////////////////////////////////////////////////////////////////////////////
/// \brief Add pythonization for TTree::SetBranchAddress.
/// \param[in] self Always null, since this is a module function.
/// \param[in] args Pointer to a Python tuple object containing the arguments
/// received from Python.
///
/// Modify the behaviour of SetBranchAddress so that proxy references can be passed
/// as arguments from the Python side, more precisely in cases where the C++
/// implementation of the method expects the address of a pointer.
///
/// For example:
/// ~~~{.py}
/// v = ROOT.std.vector('int')()
/// t.SetBranchAddress("my_vector_branch", v)
/// ~~~
PyObject *PyROOT::SetBranchAddressPyz(PyObject * /* self */, PyObject *args)
{
PyObject *treeObj = nullptr, *name = nullptr, *address = nullptr;

int argc = PyTuple_GET_SIZE(args);

// Look for the (const char*, void*) overload
auto argParseStr = "OUO:SetBranchAddress";
if (argc == 3 && PyArg_ParseTuple(args, argParseStr, &treeObj, &name, &address)) {

auto tree = (TTree *)GetTClass(treeObj)->DynamicCast(TTree::Class(), CPyCppyy::Instance_AsVoidPtr(treeObj));

if (!tree) {
PyErr_SetString(PyExc_TypeError,
"TTree::SetBranchAddress must be called with a TTree instance as first argument");
return nullptr;
}

auto branchName = PyUnicode_AsUTF8(name);
auto branch = tree->GetBranch(branchName);
if (!branch) {
PyErr_SetString(PyExc_TypeError, "TTree::SetBranchAddress must be called with a valid branch name");
return nullptr;
}

bool isLeafList = branch->IsA() == TBranch::Class();

void *buf = 0;
if (CPyCppyy::Instance_Check(address)) {
((CPPInstance *)address)->GetDatamemberCache(); // force creation of cache

if (((CPPInstance *)address)->fFlags & CPPInstance::kIsReference || isLeafList)
buf = CPyCppyy::Instance_AsVoidPtr(address);
else
buf = (void *)&(((CPPInstance *)address)->GetObjectRaw());
} else
Utility::GetBuffer(address, '*', 1, buf, false);

if (buf != nullptr) {
auto res = tree->SetBranchAddress(PyUnicode_AsUTF8(name), buf);
return PyInt_FromLong(res);
}
}

// Not the overload we wanted to pythonize, return None
Py_RETURN_NONE;
}

////////////////////////////////////////////////////////////////////////////
/// Try to match the arguments of TTree::Branch to the following overload:
/// - ( const char*, void*, const char*, Int_t = 32000 )
Expand Down

0 comments on commit 5b0c854

Please sign in to comment.