PyROOTPickle.cpp

Go to the documentation of this file.

/***********************************************************************************\
* (c) Copyright 1998-2024 CERN for the benefit of the LHCb and ATLAS collaborations *
*                                                                                   *
* This software is distributed under the terms of the Apache version 2 licence,     *
* copied verbatim in the file "LICENSE".                                            *
*                                                                                   *
* In applying this licence, CERN does not waive the privileges and immunities       *
* granted to it by virtue of its status as an Intergovernmental Organization        *
* or submit itself to any jurisdiction.                                             *
\***********************************************************************************/
#ifdef __ICC
// disable icc remark #2259: non-pointer conversion from "X" to "Y" may lose significant bits
//   TODO: To be removed, since it comes from ROOT
#  pragma warning( disable : 2259 )
#endif
 
#include <GaudiMP/PyROOTPickle.h>
#include <RVersion.h>
#include <TBufferFile.h>
#include <TClass.h>
#include <TClassRef.h>
#include <TPython.h>
 
//- data _______________________________________________________________________
#if ROOT_VERSION_CODE < ROOT_VERSION( 5, 19, 0 )
static PyObject* gExpand = 0;
#endif
 
namespace GaudiMP {
 
#if ROOT_VERSION_CODE < ROOT_VERSION( 5, 19, 0 )
 
  PyObject* ObjectProxyReduce( PyObject* self ) {
    // Turn the object proxy instance into a character stream and return for
    // pickle, together with the callable object that can restore the stream
    // into the object proxy instance.
 
    void* vself = TPython::ObjectProxy_AsVoidPtr( self ); // checks type
    if ( !vself ) {
      PyErr_SetString( PyExc_TypeError, "__reduce__ requires an object proxy instance as first argument" );
      return 0;
    }
 
    PyObject* nattr  = PyObject_GetAttrString( (PyObject*)self->ob_type, (char*)"__name__" );
    PyObject* pyname = PyObject_Str( nattr );
    Py_DECREF( nattr );
 
    static TClass* bufferclass = TClass::GetClass( "TBufferFile" );
    TClass*        klass       = TClass::GetClass( PyString_AS_STRING( pyname ) );
 
    // no cast is needed, but WriteObject taking a TClass argument is protected,
    // so use WriteObjectAny()
    TBufferFile* buf = 0;
    if ( klass == bufferclass ) {
      buf = (TBufferFile*)vself;
    } else {
      static TBufferFile buffer( TBuffer::kWrite );
      buffer.Reset();
      if ( buffer.WriteObjectAny( vself, klass ) != 1 ) {
        PyErr_Format( PyExc_IOError, "could not stream object of type %s", PyString_AS_STRING( pyname ) );
        Py_DECREF( pyname );
        return 0;
      }
      buf = &buffer;
    }
 
    // use a string for the serialized result, as a python buffer will not copy
    // the buffer contents; use a string for the class name, used when casting
    // on reading back in
    PyObject* res2 = PyTuple_New( 2 );
    PyTuple_SET_ITEM( res2, 0, PyString_FromStringAndSize( buf->Buffer(), buf->Length() ) );
    PyTuple_SET_ITEM( res2, 1, pyname );
 
    PyObject* result = PyTuple_New( 2 );
    Py_INCREF( gExpand );
    PyTuple_SET_ITEM( result, 0, gExpand );
    PyTuple_SET_ITEM( result, 1, res2 );
 
    return result;
  }
 
  class ObjectProxy {
  public:
    enum EFlags { kNone = 0x0, kIsOwner = 0x0001, kIsReference = 0x0002 };
 
  public:
    void HoldOn() { fFlags |= kIsOwner; }
    void Release() { fFlags &= ~kIsOwner; }
 
  public: // public, as the python C-API works with C structs
    PyObject_HEAD void* fObject;
    TClassRef           fClass;
    int                 fFlags;
 
  private: // private, as the python C-API will handle creation
    ObjectProxy() {}
  };
 
  PyObject* ObjectProxyExpand( PyObject*, PyObject* args ) {
    // This method is a helper for (un)pickling of ObjectProxy instances.
    PyObject*   pybuf  = 0;
    const char* clname = 0;
    if ( !PyArg_ParseTuple( args, const_cast<char*>( "O!s:__expand__" ), &PyString_Type, &pybuf, &clname ) ) return 0;
 
    // use the PyString macro's to by-pass error checking; do not adopt the buffer,
    // as the local TBufferFile can go out of scope (there is no copying)
    void* result;
    if ( strcmp( clname, "TBufferFile" ) == 0 ) {
      TBufferFile* buf = new TBufferFile( TBuffer::kWrite );
      buf->WriteFastArray( PyString_AS_STRING( pybuf ), PyString_GET_SIZE( pybuf ) );
      result = buf;
    } else {
      TBufferFile buf( TBuffer::kRead, PyString_GET_SIZE( pybuf ), PyString_AS_STRING( pybuf ), kFALSE );
      result = buf.ReadObjectAny( 0 );
    }
    PyObject* pobj = TPython::ObjectProxy_FromVoidPtr( result, clname );
    // set Ownership of the returned object
    ObjectProxy* obj = (ObjectProxy*)pobj;
    obj->HoldOn();
    return pobj;
  }
 
  void PyROOTPickle::Initialize( PyObject* libpyroot_pymodule, PyObject* objectproxy_pytype ) {
    Py_INCREF( libpyroot_pymodule );
    PyTypeObject* pytype = (PyTypeObject*)objectproxy_pytype;
 
    static PyMethodDef s_pdefExp = { (char*)"_ObjectProxy__expand__", (PyCFunction)ObjectProxyExpand, METH_VARARGS,
                                     (char*)"internal function" };
 
    PyObject* pymname = PyString_FromString( PyModule_GetName( libpyroot_pymodule ) );
    gExpand           = PyCFunction_NewEx( &s_pdefExp, NULL, pymname );
    Py_DECREF( pymname );
    Bool_t isOk = PyObject_SetAttrString( libpyroot_pymodule, s_pdefExp.ml_name, gExpand ) == 0;
    Py_DECREF( gExpand ); // is moderately risky, but Weakref not allowed (?)
 
    if ( !isOk ) {
      Py_DECREF( libpyroot_pymodule );
      PyErr_SetString( PyExc_TypeError, "could not add expand function to libPyROOT" );
      return;
    }
 
    static PyMethodDef s_pdefRed = { (char*)"__reduce__", (PyCFunction)ObjectProxyReduce, METH_NOARGS,
                                     (char*)"internal function" };
 
    PyObject* descr = PyDescr_NewMethod( pytype, &s_pdefRed );
    isOk            = PyDict_SetItemString( pytype->tp_dict, s_pdefRed.ml_name, descr ) == 0;
    Py_DECREF( descr );
    if ( !isOk ) {
      Py_DECREF( libpyroot_pymodule );
      PyErr_SetString( PyExc_TypeError, "could not add __reduce__ function to ObjectProxy" );
      return;
    }
 
    Py_DECREF( libpyroot_pymodule );
  }
 
#else //  ROOT_VERSION_CODE < ROOT_VERSION(5,19,0)
 
  void PyROOTPickle::Initialize( PyObject*, PyObject* ) { /* dummy. It is not needed for this version of ROOT */
  }
 
#endif //  ROOT_VERSION_CODE < ROOT_VERSION(5,19,0)
 
} // namespace GaudiMP