214 lines
9 KiB
Diff
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
|
|
|