[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.

Author: guitargeek

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py

@@ -130,7 +130,7 @@
 */
 '''
 
-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:

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,

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);
 

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 )