root/root-python-3.13-TTree-SetBranchAddress-pythonization.patch
2024-06-10 22:42:59 +02:00

214 lines
9 KiB
Diff

From 0cc7d61c6f732e701ef6ed263344105b01ef9613 Mon Sep 17 00:00:00 2001
From: Jonas Rembser <jonas.rembser@cern.ch>
Date: Mon, 10 Jun 2024 12:32:42 +0200
Subject: [PATCH] [PyROOT] Implement TTree SetBranchAddress pythonization in
Python
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.
---
.../python/ROOT/_pythonization/_ttree.py | 67 ++++++++++++++++---
.../pythonizations/src/PyROOTModule.cxx | 2 -
.../pythonizations/src/PyROOTPythonize.h | 1 -
.../pyroot/pythonizations/src/TTreePyz.cxx | 63 -----------------
4 files changed, 59 insertions(+), 74 deletions(-)
diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py
index 2781189e45..18840bb43e 100644
--- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py
+++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py
@@ -130,7 +130,7 @@ ds.SetBranchAddress('structb', ms)
*/
'''
-from libROOTPythonizations import GetBranchAttr, SetBranchAddressPyz, BranchPyz
+from libROOTPythonizations import GetBranchAttr, BranchPyz
from . import pythonization
# TTree iterator
@@ -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:
diff --git a/bindings/pyroot/pythonizations/src/PyROOTModule.cxx b/bindings/pyroot/pythonizations/src/PyROOTModule.cxx
index e943e0ad4b..b74bae358e 100644
--- a/bindings/pyroot/pythonizations/src/PyROOTModule.cxx
+++ b/bindings/pyroot/pythonizations/src/PyROOTModule.cxx
@@ -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,
diff --git a/bindings/pyroot/pythonizations/src/PyROOTPythonize.h b/bindings/pyroot/pythonizations/src/PyROOTPythonize.h
index 9f2d4506b6..684f256596 100644
--- a/bindings/pyroot/pythonizations/src/PyROOTPythonize.h
+++ b/bindings/pyroot/pythonizations/src/PyROOTPythonize.h
@@ -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);
diff --git a/bindings/pyroot/pythonizations/src/TTreePyz.cxx b/bindings/pyroot/pythonizations/src/TTreePyz.cxx
index d6763fdb25..f06811c43c 100644
--- a/bindings/pyroot/pythonizations/src/TTreePyz.cxx
+++ b/bindings/pyroot/pythonizations/src/TTreePyz.cxx
@@ -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 )
--
2.45.2