Logo ROOT  
Reference Guide
TTreePyz.cxx
Go to the documentation of this file.
1// Author: Enric Tejedor CERN 06/2018
2// Original PyROOT code by Wim Lavrijsen, LBL
3
4/*************************************************************************
5 * Copyright (C) 1995-2018, Rene Brun and Fons Rademakers. *
6 * All rights reserved. *
7 * *
8 * For the licensing terms see $ROOTSYS/LICENSE. *
9 * For the list of contributors see $ROOTSYS/README/CREDITS. *
10 *************************************************************************/
11
12// Bindings
13#include "CPyCppyy.h"
14#include "PyROOTPythonize.h"
15#include "CPPInstance.h"
16#include "ProxyWrappers.h"
17#include "Converters.h"
18#include "Utility.h"
19#include "PyzCppHelpers.hxx"
20
21// ROOT
22#include "TClass.h"
23#include "TTree.h"
24#include "TBranch.h"
25#include "TBranchElement.h"
26#include "TBranchObject.h"
27#include "TLeaf.h"
28#include "TLeafElement.h"
29#include "TLeafObject.h"
30#include "TStreamerElement.h"
31#include "TStreamerInfo.h"
32
33using namespace CPyCppyy;
34
35static TBranch *SearchForBranch(TTree *tree, const char *name)
36{
37 TBranch *branch = tree->GetBranch(name);
38 if (!branch) {
39 // for benefit of naming of sub-branches, the actual name may have a trailing '.'
40 branch = tree->GetBranch((std::string(name) + '.').c_str());
41 }
42 return branch;
43}
44
45static TLeaf *SearchForLeaf(TTree *tree, const char *name, TBranch *branch)
46{
47 TLeaf *leaf = tree->GetLeaf(name);
48 if (branch && !leaf) {
49 leaf = branch->GetLeaf(name);
50 if (!leaf) {
51 TObjArray *leaves = branch->GetListOfLeaves();
52 if (leaves->GetSize() && (leaves->First() == leaves->Last())) {
53 // i.e., if unambiguously only this one
54 leaf = (TLeaf *)leaves->At(0);
55 }
56 }
57 }
58 return leaf;
59}
60
61static PyObject *BindBranchToProxy(TTree *tree, const char *name, TBranch *branch)
62{
63 // for partial return of a split object
64 if (branch->InheritsFrom(TBranchElement::Class())) {
65 TBranchElement *be = (TBranchElement *)branch;
66 if (be->GetCurrentClass() && (be->GetCurrentClass() != be->GetTargetClass()) && (0 <= be->GetID())) {
67 Long_t offset = ((TStreamerElement *)be->GetInfo()->GetElements()->At(be->GetID()))->GetOffset();
69 }
70 }
71
72 // for return of a full object
73 if (branch->IsA() == TBranchElement::Class() || branch->IsA() == TBranchObject::Class()) {
74 TClass *klass = TClass::GetClass(branch->GetClassName());
75 if (klass && branch->GetAddress())
76 return BindCppObjectNoCast(*(void **)branch->GetAddress(), Cppyy::GetScope(branch->GetClassName()));
77
78 // try leaf, otherwise indicate failure by returning a typed null-object
79 TObjArray *leaves = branch->GetListOfLeaves();
80 if (klass && !tree->GetLeaf(name) && !(leaves->GetSize() && (leaves->First() == leaves->Last())))
81 return BindCppObjectNoCast(nullptr, Cppyy::GetScope(branch->GetClassName()));
82 }
83
84 return nullptr;
85}
86
87static PyObject *WrapLeaf(TLeaf *leaf)
88{
89 if (1 < leaf->GetLenStatic() || leaf->GetLeafCount()) {
90 // array types
91 dim_t dims[] = { 1, leaf->GetNdata() }; // first entry is the number of dims
92 std::string typeName = leaf->GetTypeName();
93 Converter *pcnv = CreateConverter(typeName + '*', dims);
94
95 void *address = 0;
96 if (leaf->GetBranch())
97 address = (void *)leaf->GetBranch()->GetAddress();
98 if (!address)
99 address = (void *)leaf->GetValuePointer();
100
101 PyObject *value = pcnv->FromMemory(&address);
103
104 return value;
105 } else if (leaf->GetValuePointer()) {
106 // value types
107 Converter *pcnv = CreateConverter(leaf->GetTypeName());
108 PyObject *value = 0;
109 if (leaf->IsA() == TLeafElement::Class() || leaf->IsA() == TLeafObject::Class())
110 value = pcnv->FromMemory((void *)*(void **)leaf->GetValuePointer());
111 else
112 value = pcnv->FromMemory((void *)leaf->GetValuePointer());
114
115 return value;
116 }
117
118 return nullptr;
119}
120
121// Allow access to branches/leaves as if they were data members
123{
124 const char *name_possibly_alias = CPyCppyy_PyText_AsString(pyname);
125 if (!name_possibly_alias)
126 return 0;
127
128 // get hold of actual tree
129 auto tree = (TTree *)GetTClass(self)->DynamicCast(TTree::Class(), self->GetObject());
130
131 if (!tree) {
132 PyErr_SetString(PyExc_ReferenceError, "attempt to access a null-pointer");
133 return 0;
134 }
135
136 // deal with possible aliasing
137 const char *name = tree->GetAlias(name_possibly_alias);
138 if (!name)
139 name = name_possibly_alias;
140
141 // search for branch first (typical for objects)
142 TBranch *branch = SearchForBranch(tree, name);
143
144 if (branch) {
145 // found a branched object, wrap its address for the object it represents
146 auto proxy = BindBranchToProxy(tree, name, branch);
147 if (proxy != nullptr)
148 return proxy;
149 }
150
151 // if not, try leaf
152 TLeaf *leaf = SearchForLeaf(tree, name, branch);
153
154 if (leaf) {
155 // found a leaf, extract value and wrap with a Python object according to its type
156 auto wrapper = WrapLeaf(leaf);
157 if (wrapper != nullptr)
158 return wrapper;
159 }
160
161 // confused
162 PyErr_Format(PyExc_AttributeError, "\'%s\' object has no attribute \'%s\'", tree->IsA()->GetName(), name);
163 return 0;
164}
165
166////////////////////////////////////////////////////////////////////////////
167/// \brief Allow branches to be accessed as attributes of a tree.
168/// \param[in] self Always null, since this is a module function.
169/// \param[in] args Pointer to a Python tuple object containing the arguments
170/// received from Python.
171///
172/// Allow access to branches/leaves as if they were Python data attributes of the tree
173/// (e.g. mytree.branch)
175{
176 PyObject *pyclass = PyTuple_GetItem(args, 0);
177 Utility::AddToClass(pyclass, "__getattr__", (PyCFunction)GetAttr, METH_O);
179}
180
181////////////////////////////////////////////////////////////////////////////
182/// \brief Add pythonization for TTree::SetBranchAddress.
183/// \param[in] self Always null, since this is a module function.
184/// \param[in] args Pointer to a Python tuple object containing the arguments
185/// received from Python.
186///
187/// Modify the behaviour of SetBranchAddress so that proxy references can be passed
188/// as arguments from the Python side, more precisely in cases where the C++
189/// implementation of the method expects the address of a pointer.
190///
191/// For example:
192/// ~~~{.python}
193/// v = ROOT.std.vector('int')()
194/// t.SetBranchAddress("my_vector_branch", v)
195/// ~~~
197{
198 PyObject *treeObj = nullptr, *name = nullptr, *address = nullptr;
199
200 int argc = PyTuple_GET_SIZE(args);
201
202// Look for the (const char*, void*) overload
203#if PY_VERSION_HEX < 0x03000000
204 auto argParseStr = "OSO:SetBranchAddress";
205#else
206 auto argParseStr = "OUO:SetBranchAddress";
207#endif
208 if (argc == 3 && PyArg_ParseTuple(args, const_cast<char *>(argParseStr), &treeObj, &name, &address)) {
209
210 auto treeProxy = (CPPInstance *)treeObj;
211 auto tree = (TTree *)GetTClass(treeProxy)->DynamicCast(TTree::Class(), treeProxy->GetObject());
212
213 if (!tree) {
214 PyErr_SetString(PyExc_TypeError,
215 "TTree::SetBranchAddress must be called with a TTree instance as first argument");
216 return nullptr;
217 }
218
219 auto branchName = CPyCppyy_PyText_AsString(name);
220 auto branch = tree->GetBranch(branchName);
221 if (!branch) {
222 PyErr_SetString(PyExc_TypeError, "TTree::SetBranchAddress must be called with a valid branch name");
223 return nullptr;
224 }
225
226 bool isLeafList = branch->IsA() == TBranch::Class();
227
228 void *buf = 0;
229 if (CPPInstance_Check(address)) {
230 if (((CPPInstance *)address)->fFlags & CPPInstance::kIsReference || isLeafList)
231 buf = (void *)((CPPInstance *)address)->fObject;
232 else
233 buf = (void *)&((CPPInstance *)address)->fObject;
234 } else
235 Utility::GetBuffer(address, '*', 1, buf, false);
236
237 if (buf != nullptr) {
238 auto res = tree->SetBranchAddress(CPyCppyy_PyText_AsString(name), buf);
239 return PyInt_FromLong(res);
240 }
241 }
242
243 // Not the overload we wanted to pythonize, return None
245}
246
247////////////////////////////////////////////////////////////////////////////
248/// Try to match the arguments of TTree::Branch to the following overload:
249/// - ( const char*, void*, const char*, Int_t = 32000 )
250/// If the match succeeds, invoke Branch on the C++ tree with the right
251/// arguments.
253{
254 PyObject *treeObj = nullptr;
255 PyObject *name = nullptr, *address = nullptr, *leaflist = nullptr, *bufsize = nullptr;
256
257 if (PyArg_ParseTuple(args, const_cast<char *>("OO!OO!|O!:Branch"),
258 &treeObj,
260 &address,
261 &CPyCppyy_PyText_Type, &leaflist,
262 &PyInt_Type, &bufsize)) {
263
264 auto treeProxy = (CPPInstance *)treeObj;
265 auto tree = (TTree *)GetTClass(treeProxy)->DynamicCast(TTree::Class(), treeProxy->GetObject());
266 if (!tree) {
267 PyErr_SetString(PyExc_TypeError, "TTree::Branch must be called with a TTree instance as first argument");
268 return nullptr;
269 }
270
271 void *buf = nullptr;
272 if (CPPInstance_Check(address))
273 buf = (void *)((CPPInstance *)address)->GetObject();
274 else
275 Utility::GetBuffer(address, '*', 1, buf, false);
276
277 if (buf) {
278 TBranch *branch = nullptr;
279 if (argc == 5) {
280 branch = tree->Branch(CPyCppyy_PyText_AsString(name), buf, CPyCppyy_PyText_AsString(leaflist),
281 PyInt_AS_LONG(bufsize));
282 } else {
283 branch = tree->Branch(CPyCppyy_PyText_AsString(name), buf, CPyCppyy_PyText_AsString(leaflist));
284 }
285
286 return BindCppObject(branch, Cppyy::GetScope("TBranch"));
287 }
288 }
289 PyErr_Clear();
290
292}
293
294////////////////////////////////////////////////////////////////////////////
295/// Try to match the arguments of TTree::Branch to one of the following
296/// overloads:
297/// - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 )
298/// - ( const char*, T**, Int_t = 32000, Int_t = 99 )
299/// If the match succeeds, invoke Branch on the C++ tree with the right
300/// arguments.
302{
303 PyObject *treeObj = nullptr;
304 PyObject *name = nullptr, *clName = nullptr, *address = nullptr, *bufsize = nullptr, *splitlevel = nullptr;
305
306 auto bIsMatch = false;
307 if (PyArg_ParseTuple(args, const_cast<char *>("OO!O!O|O!O!:Branch"),
308 &treeObj,
310 &CPyCppyy_PyText_Type, &clName,
311 &address,
312 &PyInt_Type, &bufsize,
313 &PyInt_Type, &splitlevel)) {
314 bIsMatch = true;
315 } else {
316 PyErr_Clear();
317 if (PyArg_ParseTuple(args, const_cast<char *>("OO!O|O!O!"),
318 &treeObj,
320 &address,
321 &PyInt_Type, &bufsize,
322 &PyInt_Type, &splitlevel)) {
323 bIsMatch = true;
324 } else {
325 PyErr_Clear();
326 }
327 }
328
329 if (bIsMatch) {
330 auto treeProxy = (CPPInstance *)treeObj;
331 auto tree = (TTree *)GetTClass(treeProxy)->DynamicCast(TTree::Class(), treeProxy->GetObject());
332 if (!tree) {
333 PyErr_SetString(PyExc_TypeError, "TTree::Branch must be called with a TTree instance as first argument");
334 return nullptr;
335 }
336
337 std::string klName = clName ? CPyCppyy_PyText_AsString(clName) : "";
338 void *buf = nullptr;
339
340 if (CPPInstance_Check(address)) {
341 if (((CPPInstance *)address)->fFlags & CPPInstance::kIsReference)
342 buf = (void *)((CPPInstance *)address)->fObject;
343 else
344 buf = (void *)&((CPPInstance *)address)->fObject;
345
346 if (!clName) {
347 klName = GetTClass((CPPInstance *)address)->GetName();
348 argc += 1;
349 }
350 } else {
351 Utility::GetBuffer(address, '*', 1, buf, false);
352 }
353
354 if (buf && !klName.empty()) {
355 TBranch *branch = nullptr;
356 if (argc == 4) {
357 branch = tree->Branch(CPyCppyy_PyText_AsString(name), klName.c_str(), buf);
358 } else if (argc == 5) {
359 branch = tree->Branch(CPyCppyy_PyText_AsString(name), klName.c_str(), buf, PyInt_AS_LONG(bufsize));
360 } else if (argc == 6) {
361 branch = tree->Branch(CPyCppyy_PyText_AsString(name), klName.c_str(), buf, PyInt_AS_LONG(bufsize),
362 PyInt_AS_LONG(splitlevel));
363 }
364
365 return BindCppObject(branch, Cppyy::GetScope("TBranch"));
366 }
367 }
368
370}
371
372////////////////////////////////////////////////////////////////////////////
373/// \brief Add pythonization for TTree::Branch.
374/// \param[in] self Always null, since this is a module function.
375/// \param[in] args Pointer to a Python tuple object containing the arguments
376/// received from Python.
377///
378/// Modify the behaviour of Branch so that proxy references can be passed
379/// as arguments from the Python side, more precisely in cases where the C++
380/// implementation of the method expects the address of a pointer.
381///
382/// For example:
383/// ~~~{.python}
384/// v = ROOT.std.vector('int')()
385/// t.Branch('my_vector_branch', v)
386/// ~~~
387///
388/// The following signatures are treated in this pythonization:
389/// - ( const char*, void*, const char*, Int_t = 32000 )
390/// - ( const char*, const char*, T**, Int_t = 32000, Int_t = 99 )
391/// - ( const char*, T**, Int_t = 32000, Int_t = 99 )
393{
394 int argc = PyTuple_GET_SIZE(args);
395
396 if (argc >= 3) { // We count the TTree proxy object too
397 auto branch = TryBranchLeafListOverload(argc, args);
398 if (branch != Py_None)
399 return branch;
400
401 branch = TryBranchPtrToPtrOverloads(argc, args);
402 if (branch != Py_None)
403 return branch;
404 }
405
406 // Not the overload we wanted to pythonize, return None
408}
#define CPyCppyy_PyText_AsString
Definition: CPyCppyy.h:97
#define Py_RETURN_NONE
Definition: CPyCppyy.h:281
#define CPyCppyy_PyText_Type
Definition: CPyCppyy.h:115
void Class()
Definition: Class.C:29
PyInt_FromLong
Definition: Converters.cxx:858
unsigned int fFlags
_object PyObject
Definition: PyMethodBase.h:41
TClass * GetTClass(const CPyCppyy::CPPInstance *pyobj)
long Long_t
Definition: RtypesCore.h:52
@ kIsReference
Definition: TDictionary.h:82
char name[80]
Definition: TGX11.cxx:109
#define pyname
Definition: TMCParticle.cxx:19
static TBranch * SearchForBranch(TTree *tree, const char *name)
Definition: TTreePyz.cxx:35
static PyObject * BindBranchToProxy(TTree *tree, const char *name, TBranch *branch)
Definition: TTreePyz.cxx:61
PyObject * TryBranchLeafListOverload(int argc, PyObject *args)
Try to match the arguments of TTree::Branch to the following overload:
Definition: TTreePyz.cxx:252
PyObject * TryBranchPtrToPtrOverloads(int argc, PyObject *args)
Try to match the arguments of TTree::Branch to one of the following overloads:
Definition: TTreePyz.cxx:301
static TLeaf * SearchForLeaf(TTree *tree, const char *name, TBranch *branch)
Definition: TTreePyz.cxx:45
PyObject * GetAttr(CPPInstance *self, PyObject *pyname)
Definition: TTreePyz.cxx:122
static PyObject * WrapLeaf(TLeaf *leaf)
Definition: TTreePyz.cxx:87
virtual PyObject * FromMemory(void *address)
Definition: Converters.cxx:422
A Branch for the case of an object.
Int_t GetID() const
TStreamerInfo * GetInfo() const
Get streamer info for the branch class.
TClass * GetCurrentClass()
Return a pointer to the current type of the data member corresponding to branch element.
virtual TClass * GetTargetClass()
char * GetObject() const
Return a pointer to our object.
A TTree is a list of TBranches.
Definition: TBranch.h:91
virtual TLeaf * GetLeaf(const char *name) const
Return pointer to the 1st Leaf named name in thisBranch.
Definition: TBranch.cxx:1919
virtual const char * GetClassName() const
Return the name of the user class whose content is stored in this branch, if any.
Definition: TBranch.cxx:1312
virtual char * GetAddress() const
Definition: TBranch.h:210
Int_t GetOffset() const
Definition: TBranch.h:233
TObjArray * GetListOfLeaves()
Definition: TBranch.h:245
TClass instances represent classes, structs and namespaces in the ROOT type system.
Definition: TClass.h:80
void * DynamicCast(const TClass *base, void *obj, Bool_t up=kTRUE)
Cast obj of this class type up to baseclass cl if up is true.
Definition: TClass.cxx:4878
static TClass * GetClass(const char *name, Bool_t load=kTRUE, Bool_t silent=kFALSE)
Static method returning pointer to TClass of the specified class name.
Definition: TClass.cxx:2948
virtual Int_t GetSize() const
Return the capacity of the collection, i.e.
Definition: TCollection.h:182
A TLeaf describes individual elements of a TBranch See TBranch structure in TTree.
Definition: TLeaf.h:49
virtual void * GetValuePointer() const
Definition: TLeaf.h:129
virtual const char * GetTypeName() const
Definition: TLeaf.h:130
virtual TLeaf * GetLeafCount() const
If this leaf stores a variable-sized array or a multi-dimensional array whose last dimension has vari...
Definition: TLeaf.h:112
virtual Int_t GetNdata() const
Definition: TLeaf.h:127
TBranch * GetBranch() const
Definition: TLeaf.h:107
virtual Int_t GetLenStatic() const
Return the fixed length of this leaf.
Definition: TLeaf.h:123
virtual const char * GetName() const
Returns name of object.
Definition: TNamed.h:47
An array of TObjects.
Definition: TObjArray.h:37
TObject * Last() const
Return the object in the last filled slot. Returns 0 if no entries.
Definition: TObjArray.cxx:506
TObject * First() const
Return the object in the first slot.
Definition: TObjArray.cxx:496
TObject * At(Int_t idx) const
Definition: TObjArray.h:166
virtual Bool_t InheritsFrom(const char *classname) const
Returns kTRUE if object inherits from class "classname".
Definition: TObject.cxx:443
TObjArray * GetElements() const
A TTree represents a columnar dataset.
Definition: TTree.h:78
Py_ssize_t GetBuffer(PyObject *pyobject, char tc, int size, void *&buf, bool check=true)
Definition: Utility.cxx:614
bool AddToClass(PyObject *pyclass, const char *label, PyCFunction cfunc, int flags=METH_VARARGS)
Definition: Utility.cxx:169
PyObject * BindCppObjectNoCast(Cppyy::TCppObject_t object, Cppyy::TCppType_t klass, const unsigned flags=0)
bool CPPInstance_Check(T *object)
Definition: CPPInstance.h:118
PyObject * BindCppObject(Cppyy::TCppObject_t object, Cppyy::TCppType_t klass, const unsigned flags=0)
CPYCPPYY_EXTERN void DestroyConverter(Converter *p)
CPYCPPYY_EXTERN Converter * CreateConverter(const std::string &name, Py_ssize_t *dims=nullptr)
RPY_EXPORTED TCppScope_t GetScope(const std::string &scope_name)
PyObject * SetBranchAddressPyz(PyObject *self, PyObject *args)
Add pythonization for TTree::SetBranchAddress.
Definition: TTreePyz.cxx:196
PyObject * BranchPyz(PyObject *self, PyObject *args)
Add pythonization for TTree::Branch.
Definition: TTreePyz.cxx:392
PyObject * AddBranchAttrSyntax(PyObject *self, PyObject *args)
Allow branches to be accessed as attributes of a tree.
Definition: TTreePyz.cxx:174
Definition: tree.py:1