diff --git a/bindings/pyroot/pythonizations/CMakeLists.txt b/bindings/pyroot/pythonizations/CMakeLists.txt index 92dfa0df5640e..34c815e257157 100644 --- a/bindings/pyroot/pythonizations/CMakeLists.txt +++ b/bindings/pyroot/pythonizations/CMakeLists.txt @@ -13,7 +13,6 @@ set(cpp_sources src/RPyROOTApplication.cxx src/GenericPyz.cxx src/TClassPyz.cxx - src/TTreePyz.cxx src/CPPInstancePyz.cxx ) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py index 77cd655577c3a..971d935975031 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py @@ -159,8 +159,6 @@ \endpythondoc """ -from ROOT.libROOTPythonizations import BranchPyz, GetBranchAttr - from . import pythonization from ._memory_utils import ( _constructor_releasing_ownership, @@ -170,6 +168,46 @@ from ._rvec import _get_cpp_type_from_numpy_type +_branch_lookups = None + + +class _BranchLookups(object): + """The entities that reading a branch needs, resolved once. + + Every resolution through the ROOT module goes through the facade's + __getattr__, which dominated the cost of `tree.branch` when done inline; + none of these can change over the lifetime of a session. + """ + + __slots__ = ("helpers", "branch_element", "branch_object", "leaf_element", "leaf_object", "ll") + + +def _lookups(): + """Return the cached branch lookups, resolving them on first use. + + The helpers in ROOT::Internal::TreeUtils exist because + TBranch::GetAddress() and TBranchElement::GetObject() return char* and so + reach Python as str, with the pointer value gone. Two threads racing here + resolve the same entities, so no lock is needed. + """ + global _branch_lookups + + if _branch_lookups is None: + import ROOT + from cppyy import ll + + lookups = _BranchLookups() + lookups.helpers = ROOT.Internal.TreeUtils + lookups.branch_element = ROOT.TBranchElement.Class() + lookups.branch_object = ROOT.TBranchObject.Class() + lookups.leaf_element = ROOT.TLeafElement.Class() + lookups.leaf_object = ROOT.TLeafObject.Class() + lookups.ll = ll + _branch_lookups = lookups + + return _branch_lookups + + # TTree iterator def _TTree__iter__(self): i = 0 @@ -305,18 +343,224 @@ def _SetBranchAddress(self, bname, addr, *args, **kwargs): return self._OriginalSetBranchAddress(bname, addr, ptr=tbranch_ptr, realClass=cl, datatype=tp, isptr=False) +def _get_address_of(obj): + """Return the address of the buffer or proxied object `obj` points at. + + Returns None for anything that does not carry an address, and deliberately + also for text-like objects, which do but are never meant as a branch buffer. + """ + import ROOT + + if isinstance(obj, ROOT._cppyy.types.Instance): + return ROOT._cppyy.addressof(instance=obj, byref=False) + + if isinstance(obj, (str, bytes)): + return None + + # A plain integer address is left to the original overloads. It must not + # reach ll.addressof, which converts it to a C long: that overflows where + # long is 32 bits (Windows) and raises OverflowError instead of TypeError. + if isinstance(obj, int): + return None + + try: + return ROOT._cppyy.ll.addressof(obj) + except TypeError: + return None + + +def _try_branch_leaf_list_overload(self, args): + """Try to match TTree::Branch(const char*, void*, const char*, Int_t = 32000).""" + if not (3 <= len(args) <= 4): + return None + name, address, leaflist = args[0], args[1], args[2] + if not isinstance(name, str) or not isinstance(leaflist, str): + return None + if len(args) == 4 and not isinstance(args[3], int): + return None + + import ctypes + + buf = _get_address_of(address) + if not buf: + return None + + return self._OriginalBranch(name, ctypes.c_void_p(buf), leaflist, *args[3:]) + + +def _try_branch_ptr_to_ptr_overloads(self, args): + """Try to match one of the TTree::Branch overloads taking a T**: + + - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 ) + - ( const char*, T**, Int_t = 32000, Int_t = 99 ) + """ + import ROOT + + if len(args) < 2 or not isinstance(args[0], str): + return None + + name = args[0] + if isinstance(args[1], str): + # the class name is given explicitly + class_name, address, rest = args[1], args[2] if len(args) > 2 else None, args[3:] + else: + class_name, address, rest = None, args[1], args[2:] + + if address is None or any(not isinstance(arg, int) for arg in rest) or len(rest) > 2: + return None + + if isinstance(address, ROOT._cppyy.types.Instance): + # Hand the proxy to a helper taking a T** and let cppyy work out which + # address that is: where it lives (the proxy itself, or the caller's + # pointer behind a reference proxy) is a rule only the bindings know. + # T is the proxied type, not class_name, which the caller is free to + # give as a base or equivalent spelling of it. + proxy_type = type(address).__cpp_name__ + brancher = _lookups().helpers.BranchPtrToPtr[proxy_type] + return brancher(self, name, class_name or proxy_type, address, *rest) + + buf = _get_address_of(address) + if not buf or not class_name: + return None + + import ctypes + + return self._OriginalBranch(name, class_name, ctypes.c_void_p(buf), *rest) + + def _Branch(self, *args): - # Modify the behaviour if args is one of: - # ( const char*, void*, const char*, Int_t = 32000 ) - # ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 ) - # ( const char*, T**, Int_t = 32000, Int_t = 99 ) - res = BranchPyz(self, *args) + """ + Pythonization for TTree::Branch. - if res is None: - # Fall back to the original implementation for the rest of overloads - res = self._OriginalBranch(*args) + Modify the behaviour of Branch 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.Branch('my_vector_branch', v) + ``` + + The following signatures are treated in this pythonization: + - ( const char*, void*, const char*, Int_t = 32000 ) + - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 ) + - ( const char*, T**, Int_t = 32000, Int_t = 99 ) + """ + if len(args) >= 2: + res = _try_branch_leaf_list_overload(self, args) + if res is not None: + return res + + res = _try_branch_ptr_to_ptr_overloads(self, args) + if res is not None: + return res + + # Fall back to the original implementation for the rest of overloads + return self._OriginalBranch(*args) + + +def _search_for_branch(tree, name): + branch = tree.GetBranch(name) + if not branch: + # for benefit of naming of sub-branches, the actual name may have a + # trailing '.' + branch = tree.GetBranch(name + ".") + return branch + + +def _has_single_leaf(branch): + leaves = branch.GetListOfLeaves() + # i.e. if unambiguously only this one + return bool(leaves.GetSize()) and leaves.First() == leaves.Last() + + +def _search_for_leaf(tree, name, branch): + leaf = tree.GetLeaf(name) + if branch and not leaf: + leaf = branch.GetLeaf(name) + if not leaf and _has_single_leaf(branch): + leaf = branch.GetListOfLeaves().At(0) + return leaf + + +def _resolve_branch(tree, name, branch): + """Return the address and type name of the object held by a branch. + + (None, "") means the branch holds no object (the caller should look for a + leaf instead); address 0 with a non-empty type name means failure, reported + to the user as a typed null object. + """ + lookups = _lookups() + helpers = lookups.helpers - return res + # for partial return of a split object + if branch.InheritsFrom(lookups.branch_element): + current_class = branch.GetCurrentClass() + if current_class and current_class != branch.GetTargetClass() and branch.GetID() >= 0: + offset = branch.GetInfo().GetElements().At(branch.GetID()).GetOffset() + return helpers.GetBranchElementObject(branch) + offset, current_class.GetName() + + # for return of a full object + if branch.IsA() in (lookups.branch_element, lookups.branch_object): + address = helpers.GetBranchAddress(branch) + if address: + # the branch address is that of the pointer to the object + return lookups.ll.bind_value("intptr_t", address), branch.GetClassName() + + # try leaf, otherwise indicate failure by returning a typed null object + if not tree.GetLeaf(name) and not _has_single_leaf(branch): + return 0, branch.GetClassName() + + return None, "" + + +def _get_multi_dims(title): + """Extract the static dimensions from the title of a TLeaf. + + The title carries them as `name[dim1][dim2]...`; TLeaf offers no other + way to get at them. + """ + import re + + return [int(dim) for dim in re.findall(r"\[([^\]]*)\]", title) if dim] + + +def _wrap_leaf(leaf): + """Read the value of a leaf for the entry the tree is currently on.""" + lookups = _lookups() + ll = lookups.ll + + if leaf.GetLenStatic() > 1 or leaf.GetLeafCount(): + # array types + is_static = leaf.GetLenStatic() > 1 + type_name = leaf.GetTypeName() + + dims = [leaf.GetNdata()] + title = leaf.GetTitle() + if title.count("[") >= 2: + # multidimensional array case + dims = _get_multi_dims(title) + + address = 0 + branch = leaf.GetBranch() + if branch: + address = lookups.helpers.GetBranchAddress(branch) + if not address: + address = ll.addressof(leaf.GetValuePointer()) + + return ll.bind_value(type_name + ("[]" if is_static else "*"), address, dims) + + value_pointer = leaf.GetValuePointer() + if value_pointer: + # value types + address = ll.addressof(value_pointer) + if leaf.IsA() in (lookups.leaf_element, lookups.leaf_object): + # the leaf holds a pointer to the value, rather than the value + address = ll.bind_value("intptr_t", address) + return ll.bind_value(leaf.GetTypeName(), address) + + return None def _TTree__getattr__(self, key): @@ -326,21 +570,35 @@ def _TTree__getattr__(self, key): Allow access to branches/leaves as if they were Python data attributes of the tree (e.g. mytree.branch). - To avoid using the CPyCppyy API, any necessary cast is done here on the - Python side. The GetBranchAttr() function encodes a necessary cast in the - second element of the output tuple, which is a string with the required - type name. - Parameters: self (TTree): The instance of the TTree object from which the attribute is being retrieved. key (str): The name of the branch to retrieve from the TTree object. """ - import ROOT - - out, cast_type = GetBranchAttr(self, key) - if cast_type: - out = ROOT._cppyy.ll.cast[cast_type](out) - return out + ll = _lookups().ll + + # deal with possible aliasing + name = self.GetAlias(key) or key + + # search for branch first (typical for objects) + branch = _search_for_branch(self, name) + + if branch: + # found a branched object, wrap its address for the object it represents + address, type_name = _resolve_branch(self, name, branch) + if type_name: + return ll.cast[type_name + "*"](address) + + # if not, try leaf + leaf = _search_for_leaf(self, name, branch) + if leaf: + # found a leaf, extract value and wrap with a Python object + # according to its type + value = _wrap_leaf(leaf) + if value is not None: + return value + + # confused + raise AttributeError("'{}' object has no attribute '{}'".format(self.IsA().GetName(), name)) def _TTree_CloneTree(self, *args, **kwargs): diff --git a/bindings/pyroot/pythonizations/src/PyROOTModule.cxx b/bindings/pyroot/pythonizations/src/PyROOTModule.cxx index d7952330e2ef6..2cdf10ce6570f 100644 --- a/bindings/pyroot/pythonizations/src/PyROOTModule.cxx +++ b/bindings/pyroot/pythonizations/src/PyROOTModule.cxx @@ -183,12 +183,8 @@ void GetBuffer(PyObject *pyobject, void *&buf) static PyMethodDef gPyROOTMethods[] = { {"AddCPPInstancePickling", (PyCFunction)PyROOT::AddCPPInstancePickling, METH_NOARGS, "Add a custom pickling mechanism for Cppyy Python proxy objects"}, - {"GetBranchAttr", (PyCFunction)PyROOT::GetBranchAttr, METH_VARARGS, - "Allow to access branches as tree attributes"}, {"AddTClassDynamicCastPyz", (PyCFunction)PyROOT::AddTClassDynamicCastPyz, METH_VARARGS, "Cast the void* returned by TClass::DynamicCast to the right type"}, - {"BranchPyz", (PyCFunction)PyROOT::BranchPyz, METH_VARARGS, - "Fully enable the use of TTree::Branch from Python"}, {"AddPrettyPrintingPyz", (PyCFunction)PyROOT::AddPrettyPrintingPyz, METH_VARARGS, "Add pretty printing pythonization"}, {"InitApplication", (PyCFunction)PyROOT::RPyROOTApplication::InitApplication, METH_VARARGS, diff --git a/bindings/pyroot/pythonizations/src/PyROOTPythonize.h b/bindings/pyroot/pythonizations/src/PyROOTPythonize.h index aec7ed708b3f9..7246bf2970062 100644 --- a/bindings/pyroot/pythonizations/src/PyROOTPythonize.h +++ b/bindings/pyroot/pythonizations/src/PyROOTPythonize.h @@ -20,8 +20,6 @@ PyObject *AddCPPInstancePickling(PyObject *self, PyObject *args); PyObject *AddPrettyPrintingPyz(PyObject *self, PyObject *args); -PyObject *GetBranchAttr(PyObject *self, PyObject *args); -PyObject *BranchPyz(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 deleted file mode 100644 index ec8bf202f6354..0000000000000 --- a/bindings/pyroot/pythonizations/src/TTreePyz.cxx +++ /dev/null @@ -1,409 +0,0 @@ -// Author: Enric Tejedor CERN 06/2018 -// Original PyROOT code by Wim Lavrijsen, LBL - -/************************************************************************* - * Copyright (C) 1995-2018, Rene Brun and Fons Rademakers. * - * All rights reserved. * - * * - * For the licensing terms see $ROOTSYS/LICENSE. * - * For the list of contributors see $ROOTSYS/README/CREDITS. * - *************************************************************************/ - -// Bindings -#include - -// TODO: refactor public cpyrt API such that this forward declaration is not -// needed anymore. Including "cpyrt/API.h" should be enough. -namespace cppjit::cpyrt { -typedef Py_ssize_t dim_t; -} // namespace cppjit::cpyrt - -#include "../../cppjit/src/interop/cppjit_interop.h" -#include "../../cppjit/src/cpyrt/CPPInstance.h" -#include "../../cppjit/src/cpyrt/ProxyWrappers.h" -#include "../../cppjit/src/cpyrt/Dimensions.h" - -#include "cpyrt/API.h" - -#include "PyROOTPythonize.h" - -// ROOT -#include "TClass.h" -#include "TTree.h" -#include "TBranch.h" -#include "TBranchElement.h" -#include "TBranchObject.h" -#include "TLeaf.h" -#include "TLeafElement.h" -#include "TLeafObject.h" -#include "TStreamerElement.h" -#include "TStreamerInfo.h" - -#include -#include - -namespace { - -// Get the TClass of the C++ object proxied by pyobj -TClass *GetTClass(PyObject *pyobj) -{ - return TClass::GetClass(cppjit::cpyrt::Instance_GetScopedFinalName(pyobj).c_str()); -} - -} // namespace - -using namespace cppjit::cpyrt; - -namespace PyROOT{ -void GetBuffer(PyObject *pyobject, void *&buf); -} - -static TBranch *SearchForBranch(TTree *tree, const char *name) -{ - TBranch *branch = tree->GetBranch(name); - if (!branch) { - // for benefit of naming of sub-branches, the actual name may have a trailing '.' - branch = tree->GetBranch((std::string(name) + '.').c_str()); - } - return branch; -} - -static TLeaf *SearchForLeaf(TTree *tree, const char *name, TBranch *branch) -{ - TLeaf *leaf = tree->GetLeaf(name); - if (branch && !leaf) { - leaf = branch->GetLeaf(name); - if (!leaf) { - TObjArray *leaves = branch->GetListOfLeaves(); - if (leaves->GetSize() && (leaves->First() == leaves->Last())) { - // i.e., if unambiguously only this one - leaf = (TLeaf *)leaves->At(0); - } - } - } - return leaf; -} - -static std::pair ResolveBranch(TTree *tree, const char *name, TBranch *branch) -{ - // for partial return of a split object - if (branch->InheritsFrom(TBranchElement::Class())) { - TBranchElement *be = (TBranchElement *)branch; - if (be->GetCurrentClass() && (be->GetCurrentClass() != be->GetTargetClass()) && (0 <= be->GetID())) { - Long_t offset = ((TStreamerElement *)be->GetInfo()->GetElements()->At(be->GetID()))->GetOffset(); - return {be->GetObject() + offset, be->GetCurrentClass()->GetName()}; - } - } - - // for return of a full object - if (branch->IsA() == TBranchElement::Class() || branch->IsA() == TBranchObject::Class()) { - if (branch->GetAddress()) - return {*(void **)branch->GetAddress(), branch->GetClassName()}; - - // try leaf, otherwise indicate failure by returning a typed null-object - TObjArray *leaves = branch->GetListOfLeaves(); - if (!tree->GetLeaf(name) && !(leaves->GetSize() && (leaves->First() == leaves->Last()))) - return {nullptr, branch->GetClassName()}; - } - - return {nullptr, ""}; -} - -/** - * @brief Extracts static dimensions from the title of a TLeaf object. - * - * The function assumes that the title of the TLeaf object contains dimensions - * in the format `[dim1][dim2]...`. - * - * @note In the current implementation of TLeaf, there is no way to extract the - * dimensions without string parsing. - * - * @param title title of the TLeaf object from which to extract dimensions. - * @return std::vector A vector containing the extracted dimensions. - */ -static std::vector getMultiDims(std::string const &title) -{ - std::vector dims; - std::stringstream ss{title}; - - while (ss.good()) { - std::string substr; - getline(ss, substr, '['); - getline(ss, substr, ']'); - if (!substr.empty()) { - dims.push_back(std::stoi(substr)); - } - } - - return dims; -} - -static PyObject *WrapLeaf(TLeaf *leaf) -{ - if (1 < leaf->GetLenStatic() || leaf->GetLeafCount()) { - bool isStatic = 1 < leaf->GetLenStatic(); - // array types - std::string typeName = leaf->GetTypeName(); - std::vector dimsVec{leaf->GetNdata()}; - std::string title = leaf->GetTitle(); - // Multidimensional array case - if (std::count(title.begin(), title.end(), '[') >= 2) { - dimsVec = getMultiDims(title); - } - cppjit::cpyrt::Dimensions dims{static_cast(dimsVec.size()), dimsVec.data()}; - Converter *pcnv = CreateConverter(typeName + (isStatic ? "[]" : "*"), dims); - - void *address = 0; - if (leaf->GetBranch()) - address = (void *)leaf->GetBranch()->GetAddress(); - if (!address) - address = (void *)leaf->GetValuePointer(); - - PyObject *value = pcnv->FromMemory(&address); - cppjit::cpyrt::DestroyConverter(pcnv); - - return value; - } else if (leaf->GetValuePointer()) { - // value types - Converter *pcnv = CreateConverter(leaf->GetTypeName()); - PyObject *value = 0; - if (leaf->IsA() == TLeafElement::Class() || leaf->IsA() == TLeafObject::Class()) - value = pcnv->FromMemory((void *)*(void **)leaf->GetValuePointer()); - else - value = pcnv->FromMemory((void *)leaf->GetValuePointer()); - cppjit::cpyrt::DestroyConverter(pcnv); - - return value; - } - - return nullptr; -} - -// Allow access to branches/leaves as if they were data members Returns a -// Python tuple where the first element is either the desired cpyrt proxy, -// or an address that still needs to be wrapped by the caller in a proxy using -// cppyy.ll.cast. In the latter case, the second tuple element is the target -// type name. Otherwise, the second element is an empty string. -PyObject *PyROOT::GetBranchAttr(PyObject * /*self*/, PyObject *args) -{ - PyObject *self = nullptr; - PyObject *pyname = nullptr; - - PyArg_ParseTuple(args, "OU:GetBranchAttr", &self, &pyname); - - const char *name_possibly_alias = PyUnicode_AsUTF8AndSize(pyname, nullptr); - if (!name_possibly_alias) - return nullptr; - - // get hold of actual tree - auto tree = (TTree *)GetTClass(self)->DynamicCast(TTree::Class(), cppjit::cpyrt::Instance_AsVoidPtr(self)); - - if (!tree) { - PyErr_SetString(PyExc_ReferenceError, "attempt to access a null-pointer"); - return 0; - } - - // deal with possible aliasing - const char *name = tree->GetAlias(name_possibly_alias); - if (!name) - name = name_possibly_alias; - - // search for branch first (typical for objects) - TBranch *branch = SearchForBranch(tree, name); - - if (branch) { - // found a branched object, wrap its address for the object it represents - const auto [finalAddressVoidPtr, finalTypeName] = ResolveBranch(tree, name, branch); - if (!finalTypeName.empty()) { - PyObject *outTuple = PyTuple_New(2); - PyTuple_SetItem(outTuple, 0, PyLong_FromLongLong((intptr_t)finalAddressVoidPtr)); - PyTuple_SetItem(outTuple, 1, PyUnicode_FromString((finalTypeName + "*").c_str())); - return outTuple; - } - } - - // if not, try leaf - if (TLeaf *leaf = SearchForLeaf(tree, name, branch)) { - // found a leaf, extract value and wrap with a Python object according to its type - auto wrapper = WrapLeaf(leaf); - if (wrapper != nullptr) { - PyObject *outTuple = PyTuple_New(2); - PyTuple_SetItem(outTuple, 0, wrapper); - PyTuple_SetItem(outTuple, 1, PyUnicode_FromString("")); - return outTuple; - } - } - - // confused - PyErr_Format(PyExc_AttributeError, "\'%s\' object has no attribute \'%s\'", tree->IsA()->GetName(), name); - return 0; -} - -//////////////////////////////////////////////////////////////////////////// -/// Try to match the arguments of TTree::Branch to the following overload: -/// - ( const char*, void*, const char*, Int_t = 32000 ) -/// If the match succeeds, invoke Branch on the C++ tree with the right -/// arguments. -PyObject *TryBranchLeafListOverload(int argc, PyObject *args) -{ - PyObject *treeObj = nullptr; - PyObject *name = nullptr, *address = nullptr, *leaflist = nullptr, *bufsize = nullptr; - - if (PyArg_ParseTuple(args, "OO!OO!|O!:Branch", &treeObj, &PyUnicode_Type, &name, &address, &PyUnicode_Type, - &leaflist, &PyLong_Type, &bufsize)) { - - auto tree = (TTree *)GetTClass(treeObj)->DynamicCast(TTree::Class(), cppjit::cpyrt::Instance_AsVoidPtr(treeObj)); - if (!tree) { - PyErr_SetString(PyExc_TypeError, "TTree::Branch must be called with a TTree instance as first argument"); - return nullptr; - } - - void *buf = nullptr; - if (cppjit::cpyrt::Instance_Check(address)) - buf = cppjit::cpyrt::Instance_AsVoidPtr(address); - else - PyROOT::GetBuffer(address, buf); - - if (buf) { - TBranch *branch = nullptr; - const char *nameString = PyUnicode_AsUTF8AndSize(name, nullptr); - if (!nameString) { - return nullptr; - } - const char *leaflistString = PyUnicode_AsUTF8AndSize(leaflist, nullptr); - if (!leaflistString) { - return nullptr; - } - if (argc == 5) { - branch = tree->Branch(nameString, buf, leaflistString, PyLong_AsLong(bufsize)); - } else { - branch = tree->Branch(nameString, buf, leaflistString); - } - - return BindCppObject(branch, cppjit::interop::GetScope("TBranch")); - } - } - PyErr_Clear(); - - Py_RETURN_NONE; -} - -//////////////////////////////////////////////////////////////////////////// -/// Try to match the arguments of TTree::Branch to one of the following -/// overloads: -/// - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 ) -/// - ( const char*, T**, Int_t = 32000, Int_t = 99 ) -/// If the match succeeds, invoke Branch on the C++ tree with the right -/// arguments. -PyObject *TryBranchPtrToPtrOverloads(int argc, PyObject *args) -{ - PyObject *treeObj = nullptr; - PyObject *name = nullptr, *clName = nullptr, *address = nullptr, *bufsize = nullptr, *splitlevel = nullptr; - - auto bIsMatch = false; - if (PyArg_ParseTuple(args, "OO!O!O|O!O!:Branch", &treeObj, &PyUnicode_Type, &name, &PyUnicode_Type, &clName, - &address, &PyLong_Type, &bufsize, &PyLong_Type, &splitlevel)) { - bIsMatch = true; - } else { - PyErr_Clear(); - if (PyArg_ParseTuple(args, "OO!O|O!O!", &treeObj, &PyUnicode_Type, &name, &address, &PyLong_Type, &bufsize, - &PyLong_Type, &splitlevel)) { - bIsMatch = true; - } else { - PyErr_Clear(); - } - } - - if (bIsMatch) { - auto tree = (TTree *)GetTClass(treeObj)->DynamicCast(TTree::Class(), cppjit::cpyrt::Instance_AsVoidPtr(treeObj)); - if (!tree) { - PyErr_SetString(PyExc_TypeError, "TTree::Branch must be called with a TTree instance as first argument"); - return nullptr; - } - - std::string klName; - if (clName) { - const char *clNameString = PyUnicode_AsUTF8AndSize(clName, nullptr); - if (!clNameString) { - return nullptr; - } - klName = clNameString; - } - void *buf = nullptr; - - if (cppjit::cpyrt::Instance_Check(address)) { - if (((CPPInstance *)address)->fFlags & CPPInstance::kIsReference) - buf = (void *)((CPPInstance *)address)->fObject; - else - buf = (void *)&((CPPInstance *)address)->fObject; - - if (!clName) { - klName = GetTClass(address)->GetName(); - argc += 1; - } - } else { - PyROOT::GetBuffer(address, buf); - } - - if (buf && !klName.empty()) { - TBranch *branch = nullptr; - const char *nameString = nullptr; - if (argc == 4 || argc == 5 || argc == 6) { - nameString = PyUnicode_AsUTF8AndSize(name, nullptr); - if (!nameString) { - return nullptr; - } - } - if (argc == 4) { - branch = tree->Branch(nameString, klName.c_str(), buf); - } else if (argc == 5) { - branch = tree->Branch(nameString, klName.c_str(), buf, PyLong_AsLong(bufsize)); - } else if (argc == 6) { - branch = tree->Branch(nameString, klName.c_str(), buf, PyLong_AsLong(bufsize), - PyLong_AsLong(splitlevel)); - } - - return BindCppObject(branch, cppjit::interop::GetScope("TBranch")); - } - } - - Py_RETURN_NONE; -} - -//////////////////////////////////////////////////////////////////////////// -/// \brief Add pythonization for TTree::Branch. -/// \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 Branch 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.Branch('my_vector_branch', v) -/// ~~~ -/// -/// The following signatures are treated in this pythonization: -/// - ( const char*, void*, const char*, Int_t = 32000 ) -/// - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 ) -/// - ( const char*, T**, Int_t = 32000, Int_t = 99 ) -PyObject *PyROOT::BranchPyz(PyObject * /* self */, PyObject *args) -{ - int argc = PyTuple_Size(args); - - if (argc >= 3) { // We count the TTree proxy object too - auto branch = TryBranchLeafListOverload(argc, args); - if (branch != Py_None) - return branch; - - branch = TryBranchPtrToPtrOverloads(argc, args); - if (branch != Py_None) - return branch; - } - - // Not the overload we wanted to pythonize, return None - Py_RETURN_NONE; -} diff --git a/bindings/pyroot/pythonizations/test/ttree_branch.py b/bindings/pyroot/pythonizations/test/ttree_branch.py index c078322b50a29..317e642c3adda 100644 --- a/bindings/pyroot/pythonizations/test/ttree_branch.py +++ b/bindings/pyroot/pythonizations/test/ttree_branch.py @@ -6,6 +6,54 @@ import numpy as np +class ConcurrentBranchAccess(unittest.TestCase): + """ + The first branch access does one-off setup, instantiating a template in + the interpreter among it, which a concurrent first access must survive. + Only exercised on a cold cache, so this class must run first + in this file (unittest collects classes alphabetically: keep the name + ahead of TTreeBranch). + """ + + def test01_concurrent_first_branch_and_attribute_access(self): + import threading + + nthreads = 4 + barrier = threading.Barrier(nthreads) + + trees = [] + objects = [] + for i in range(nthreads): + t = ROOT.TTree('concurrent%d' % i, 'concurrent%d' % i) + SetOwnership(t, False) + obj = ROOT.TNamed('n%d' % i, 'n%d' % i) + trees.append(t) + objects.append(obj) + + # Branch(name, object) drives the first instantiation of + # BranchPtrToPtr, the branch attribute read the one-off branch-lookup + # setup; all threads race to do both first + errors = [] + + def work(i): + try: + t, obj = trees[i], objects[i] + barrier.wait() + t.Branch('obj', obj) + t.Fill() + assert t.obj.GetName() == 'n%d' % i + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=work, args=(i,)) for i in range(nthreads)] + for th in threads: + th.start() + for th in threads: + th.join() + + self.assertFalse(errors, errors) + + class TTreeBranch(unittest.TestCase): """ Test for the pythonization of TTree::Branch, which allows to pass proxy @@ -30,6 +78,23 @@ def setUpClass(cls): }; """) + # Declared separately: TreeHelper.h, which ttree.py uses, declares an + # identical MyStruct, so the block above is rejected whole when the two + # test files share an interpreter. Keep what is only needed here out of it. + ROOT.gInterpreter.Declare(""" + #include + #include + + // Reading a pointer data member gives a proxy for a reference to a + // pointer, which is bound differently from a proxy for an object. + struct MyHolder { + std::vector *myvec = new std::vector(); + }; + + intptr_t AddressOfMyVec(MyHolder *h) { return reinterpret_cast(&h->myvec); } + intptr_t AddressOfBranch(TBranch *b) { return reinterpret_cast(b->GetAddress()); } + """) + # Helpers def create_file_and_tree(self): f = ROOT.TFile(self.filename, 'RECREATE') @@ -169,7 +234,42 @@ def test12_read_vector_branch(self): for elem in v: self.assertEqual(elem, self.fval) - def test13_write_fallback_case(self): + def test13_write_reference_proxy_branch(self): + # A proxy for a reference to a pointer, such as the one obtained by + # reading a pointer data member, holds the address of that pointer, + # while a proxy for an object holds the object itself. Branch needs the + # former in both cases; taking the latter binds the branch to the + # proxy's own memory, so that filling writes nothing and the proxy, + # a temporary here, is gone by the time the tree is filled. + f,t = self.create_file_and_tree() + + h = ROOT.MyHolder() + h.myvec.assign(self.arraysize, self.fval) + + # Assert on the address before filling: filling through a branch bound + # to a dead proxy is undefined, and would take the test down with it + for i, args in enumerate([('refvectorb0', h.myvec), + ('refvectorb1', h.myvec, 32000), + ('refvectorb2', h.myvec, 32000, 99), + ('refvectorb3', 'std::vector', h.myvec), + ('refvectorb4', 'std::vector', h.myvec, 32000), + ('refvectorb5', 'std::vector', h.myvec, 32000, 99)]): + b = t.Branch(*args) + self.assertEqual(ROOT.AddressOfBranch(b), ROOT.AddressOfMyVec(h), + 'branch {} not bound to &MyHolder::myvec'.format(i)) + + self.fill_and_close(f, t) + + def test14_read_reference_proxy_branch(self): + f,t = self.get_tree() + + for entry in t: + for v in [ getattr(entry, 'refvectorb' + str(i)) for i in range(6) ]: + self.assertEqual(len(v), self.arraysize) + for elem in v: + self.assertEqual(elem, self.fval) + + def test15_write_fallback_case(self): f,t = self.create_file_and_tree() # Test an overload that uses the original Branch proxy @@ -182,7 +282,7 @@ def test13_write_fallback_case(self): self.fill_and_close(f, t) - def test14_read_fallback_case(self): + def test16_read_fallback_case(self): f,t = self.get_tree() for entry in t: diff --git a/tree/tree/inc/TTree.h b/tree/tree/inc/TTree.h index d6fce55392fc7..0a540b86142ca 100644 --- a/tree/tree/inc/TTree.h +++ b/tree/tree/inc/TTree.h @@ -68,6 +68,7 @@ class TPrincipal; class TFriendElement; class TCut; class TVirtualIndex; +class TBranchElement; class TBranchRef; class TBasket; class TStreamerInfo; @@ -84,6 +85,11 @@ TBranch *CallBranchImpRef(TTree &tree, const char *branchname, TClass *ptrClass, Int_t bufsize = 32000, Int_t splitlevel = 99); TBranch *CallBranchImp(TTree &tree, const char *branchname, TClass *ptrClass, void *addobj, Int_t bufsize = 32000, Int_t splitlevel = 99); + +// For the Python bindings, which see the char* that TBranch::GetAddress() and +// TBranchElement::GetObject() return as a string. +Longptr_t GetBranchAddress(TBranch *branch); +Longptr_t GetBranchElementObject(TBranchElement *branch); } class TTree : public TNamed, public TAttLine, public TAttFill, public TAttMarker { @@ -765,6 +771,17 @@ class TTree : public TNamed, public TAttLine, public TAttFill, public TAttMarker ClassDefOverride(TTree, 20) // Tree descriptor (the main ROOT I/O class) }; +namespace ROOT::Internal::TreeUtils { +/// For the Python bindings: call the TTree::Branch() overload that takes the +/// address of a pointer, with that address worked out by cppyy from a T**. +template +TBranch *BranchPtrToPtr(TTree &tree, const char *name, const char *className, T **obj, Int_t bufsize = 32000, + Int_t splitlevel = 99) +{ + return tree.Branch(name, className, reinterpret_cast(obj), bufsize, splitlevel); +} +} // namespace ROOT::Internal::TreeUtils + ////////////////////////////////////////////////////////////////////////// // // // TTreeFriendLeafIter // diff --git a/tree/tree/src/TTree.cxx b/tree/tree/src/TTree.cxx index 5d6f17bc7077b..303bb3e111506 100644 --- a/tree/tree/src/TTree.cxx +++ b/tree/tree/src/TTree.cxx @@ -10296,3 +10296,13 @@ TBranch *ROOT::Internal::TreeUtils::CallBranchImp(TTree &tree, const char *branc { return tree.BranchImp(branchname, ptrClass, addobj, bufsize, splitlevel); } + +Longptr_t ROOT::Internal::TreeUtils::GetBranchAddress(TBranch *branch) +{ + return reinterpret_cast(branch->GetAddress()); +} + +Longptr_t ROOT::Internal::TreeUtils::GetBranchElementObject(TBranchElement *branch) +{ + return reinterpret_cast(branch->GetObject()); +}