PyROOTTTreePatch.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
*/
// Called from python, so only excuted single-threaded (GIL).
#include "CxxUtils/checker_macros.h"
ATLAS_NO_CHECK_FILE_THREAD_SAFETY;
 
#include "RootUtils/PyROOTTTreePatch.h"
#include "RootUtils/PyAthenaGILStateEnsure.h"
#include "Utility.h"
 
#include "Python.h"
#include "TTree.h"
#include "TBranch.h"
#include "TPython.h"
#include "TNotifyLink.h"
 
 
//***************************************************************************
// Most of the PyROOT headers are not exported.
// So we need to copy pieces of some PyROOT declarations here.
//
 
namespace PyROOT {
 
class CPPInstance {
public:
  enum EFlags { kIsReference = 0x0008 };
 
  PyObject_HEAD
  void*     fObject;
  int       fFlags;
};
 
} // namespace PyROOT
 
using namespace PyROOT;
 
 
//=========================================================================
 
 
 
namespace RootUtils {
 
 
// Interned string constants that we use for dictionary lookups.
static PyObject* notifier_str = 0;
static PyObject* pynotify_str = 0;
 
 
class TreeNotifier
  : public TNotifyLinkBase
{
public:
  TreeNotifier (TTree* tree, PyObject* treeobj_ref);
 
  virtual ~TreeNotifier();
 
  virtual Bool_t Notify() override;
 
  virtual void Clear (Option_t* /*option*/ = "") override;
 
 
private:
  PyObject* m_treeobj_ref;
 
  TObject* m_chain;
 
  TTree* m_tree;
};
 
 
TreeNotifier::TreeNotifier (TTree* tree,
                            PyObject* treeobj_ref)
  : m_treeobj_ref (treeobj_ref),
    m_chain (tree->GetNotify()),
    m_tree (tree)
{
  // We acquire ownership of treeobj_ref.
}
 
 
TreeNotifier::~TreeNotifier()
{
  RootUtils::PyGILStateEnsure gil;
  if (m_tree) {
    // Clear the reference from the tree to this object.
    m_tree->SetNotify (m_chain);
  }
  Py_XDECREF (m_treeobj_ref);
}
 
 
Bool_t TreeNotifier::Notify()
{
  RootUtils::PyGILStateEnsure gil;
 
  // Intern __pynotify__ if needed.
  if (pynotify_str == 0)
    pynotify_str = PyUnicode_InternFromString("__pynotify__");
 
  // Look for a notification object.
  PyObject* treeobj = PyWeakref_GetObject (m_treeobj_ref);
  if (treeobj) {
    PyObject** dictptr = _PyObject_GetDictPtr (treeobj);
    if (dictptr && *dictptr) {
      PyObject* notobj = PyObject_GetItem (*dictptr, pynotify_str);
      if (notobj) {
        // Got it --- call @c Notify.
        PyObject* ret =
          PyObject_CallMethod (notobj, const_cast<char*> ("Notify"), NULL);
        if (!ret) return 0;
        Py_DECREF (ret);
      }
      else
        PyErr_Clear();
    }
  }
 
  if (m_chain) m_chain->Notify();
  return true;
}
 
 
void TreeNotifier::Clear (Option_t* /*option = ""*/)
{
  m_tree = nullptr;
}
 
 
PyObject* treeSetNotify (PyObject*, PyObject* args)
{
  // Decode the objects --- the tree and the notifier object.
  PyObject* self = 0;
  PyObject* obj = 0;
  if ( ! PyArg_ParseTuple( args, const_cast< char* >( "OO:SetNotify" ),
                           &self, &obj ) )
    return 0;
 
  // Intern strings if needed.
  if (pynotify_str == 0) {
    pynotify_str = PyUnicode_InternFromString("__pynotify__");
  }
  if (notifier_str == 0) {
    notifier_str = PyUnicode_InternFromString("__notifier__");
  }
 
  // Set up notifier.
  if (!PyObject_HasAttr (self, notifier_str)) {
    // get hold of actual tree
    void* vp = TPython::CPPInstance_AsVoidPtr (self);
    TTree* tree =
      (TTree*)objectIsA(self)->DynamicCast ( TTree::Class(), vp );
 
    PyObject* treeobj_ref = PyWeakref_NewRef (self, 0);
    if (!treeobj_ref)
      return 0;
    TObject* notifier = new TreeNotifier (tree, treeobj_ref);
    PyObject* notobj = TPython::CPPInstance_FromVoidPtr (notifier, "TObject");
    setOwnership (notobj, true);
    int stat = PyObject_SetAttr (self, notifier_str, notobj);
    Py_XDECREF (notobj);
    if (stat < 0)
      return 0;
    tree->SetNotify (notifier);
  }
 
  // Install the object.
  int stat = PyObject_SetAttr (self, pynotify_str, obj);
  if (stat < 0)
    return 0;
 
  Py_INCREF (Py_None);
  return Py_None;
}
 
 
PyObject* treeGetNotify (PyObject*, PyObject* args)
{
  // Decode the objects --- the tree.
  PyObject* self = 0;
  if ( ! PyArg_ParseTuple( args, const_cast< char* >( "O:SetNotify" ),
                           &self ) )
    return 0;
 
  // Intern string if needed.
  if (pynotify_str == 0)
    pynotify_str = PyUnicode_InternFromString("__pynotify__");
 
  // Retrieve the notifier.
  PyObject* ret = PyObject_GetAttr (self, pynotify_str);
  if (!ret) {
    PyErr_Clear();
    Py_INCREF (Py_None);
    ret = Py_None;
  }
  return ret;
}
 
 
PyObject* branchSetAddress (PyObject*, PyObject* args)
{
  // Decode arguments --- the branch and the buffer.
  PyObject* self = 0;
  PyObject* address = 0;
  if ( ! PyArg_ParseTuple( args, const_cast< char* >( "OO:SetBranchAddress" ),
                           &self, &address ) )
    return 0;
 
  // The branch as a Root object.
  TBranch* branch = 0;
  if (TPython::CPPInstance_Check (self)) {
    branch =  (TBranch*)objectIsA(self)->DynamicCast
      ( TBranch::Class(),
        TPython::CPPInstance_AsVoidPtr(self) );
  }
 
  if ( ! branch ) {
    PyErr_SetString( PyExc_TypeError,
                     "TBranch::SetAddress must be called with a "
                     "TBranch instance as first argument" );
    return 0;
  }
 
  // Convert the buffer argument to an address.
  void* buf = 0;
  if ( TPython::CPPInstance_Check( address ) ) {
    if ( ((CPPInstance*)address)->fFlags & CPPInstance::kIsReference )
      buf = (void*)((CPPInstance*)address)->fObject;
    else
      buf = (void*)&((CPPInstance*)address)->fObject;
  } else
    RootUtils::GetBuffer( address, '*', 1, buf, kFALSE );
 
  // Make the call and return.
  if ( buf != 0 )
    branch->SetAddress( buf );
 
  Py_INCREF( Py_None );
  return Py_None;
}
 
 
void installMethod (PyObject* pyclass,
                    PyMethodDef& pdef, 
                    const char* name,
                    PyCFunction cfunc)
{
  pdef.ml_name = const_cast<char*> (name);
  pdef.ml_meth = cfunc;
  pdef.ml_flags = METH_VARARGS;
  pdef.ml_doc = 0;
 
  PyObject* func = PyCFunction_New (&pdef, 0);
  PyObject* method = PyInstanceMethod_New (func);
  Bool_t isOK = PyObject_SetAttrString (pyclass, pdef.ml_name, method) == 0;
 
  if (PyErr_Occurred())
    fprintf (stderr, "Py error");
  else if (!isOK) {
    fprintf (stderr, "Could not add method %s\n", name);
  }
}
 
 
void PyROOTTTreePatch::Initialize (PyObject* tree_pyclass,
                                   PyObject* /*chain_pyclass*/,
                                   PyObject* branch_pyclass)
{
# define INSTALL_METHOD(pyclass, name, func) do {  \
  static PyMethodDef pdef;                     \
  installMethod (pyclass, pdef, name, (PyCFunction)func); \
  } while(0)
 
  INSTALL_METHOD (tree_pyclass, "SetNotify",   treeSetNotify);
  INSTALL_METHOD (tree_pyclass, "GetNotify",   treeGetNotify);
  INSTALL_METHOD (branch_pyclass, "SetAddress",   branchSetAddress);
#undef INSTALL_METHOD
}
 
 
} // namespace RootUtils