ROOT logo
// Author: Wim Lavrijsen   March 2008

// Bindings
#include "PyROOT.h"
#include "TPySelector.h"
#include "TPyReturn.h"
#include "ObjectProxy.h"
#include "MethodProxy.h"
#include "RootWrapper.h"
#include "MemoryRegulator.h"

//- ROOT
#include "TPython.h"
#include "TString.h"

//______________________________________________________________________________
//                      Python equivalent PROOF base class
//                      ==================================
//
// The problem with deriving a python class from a PyROOT bound class and then
// handing it back to a C++ framework, is that the virtual function dispatching
// of C++ is completely oblivious to the methods overridden in python. To work
// within the PROOF (C++) framework, a python class should derive from the class
// TPySelector. This class provides the proper overrides on the C++ side, and
// then forwards them, as apropriate, to python.
//
// This is an example set of scripts:
//
// ### PROOF running script, very close to equivalent .C (prooftest.py)
// import time
// from ROOT import *
//
// dataset = TDSet( 'TTree', 'h42' )
// dataset.Add( 'root:// .... myfile.root' )
//
// proof = TProof.Open('')
// time.sleep(1)                     # needed for GUI to settle
// print dataset.Process( 'TPySelector', 'aapje' )
// ### EOF
//
// ### selector module (aapje.py, name has to match as per above)
// from ROOT import TPySelector
//
// class MyPySelector( TPySelector ):
//    def Begin( self ):
//       print 'py: beginning'
//
//    def SlaveBegin( self, tree ):
//       print 'py: slave beginning'
//
//    def Process( self, entry ):
//       if self.fChain.GetEntry( entry ) <= 0:
//          return 0
//       print 'py: processing', self.fChain.MyVar
//       return 1
//
//   def SlaveTerminate( self ):
//       print 'py: slave terminating'
//
//   def Terminate( self ):
//       print 'py: terminating'
// ### EOF


//- data ---------------------------------------------------------------------
ClassImp(TPySelector)


//- private helpers ----------------------------------------------------------
void TPySelector::SetupPySelf()
{
   if ( fPySelf && fPySelf != Py_None )
      return;                      // already created ...

   TString impst = TString::Format( "import %s", GetOption() );

// use TPython to ensure that the interpreter is initialized
   if ( ! TPython::Exec( (const char*)impst ) ) {
      Abort( "failed to load provided python module" );  // Exec already printed error trace
      return;
   }

// get the TPySelector python class
   PyObject* tpysel = PyObject_GetAttrString(
      PyImport_AddModule( const_cast< char* >( "libPyROOT" ) ),
      const_cast< char* >( "TPySelector" ) );

// get handle to the module
   PyObject* pymod = PyImport_AddModule( const_cast< char* >( GetOption() ) );

// get the module dictionary to loop over
   PyObject* dict = PyModule_GetDict( pymod );
   Py_INCREF( dict );

// locate the TSelector derived class
   PyObject* allvalues = PyDict_Values( dict );

   PyObject* pyclass = 0;
   for ( int i = 0; i < PyList_GET_SIZE( allvalues ); ++i ) {
      PyObject* value = PyList_GET_ITEM( allvalues, i );
      Py_INCREF( value );

      if ( PyType_Check( value ) && PyObject_IsSubclass( value, tpysel ) ) {
         if ( PyObject_Compare(	value, tpysel ) ) {    // i.e., if not equal
            pyclass = value;
            break;
         }
      }

      Py_DECREF( value );
   }

   Py_DECREF( allvalues );
   Py_DECREF( dict );
   Py_DECREF( tpysel );

   if ( ! pyclass ) {
      Abort( "no TSelector derived class available in provided module" );
      return;
   }

   PyObject* args = PyTuple_New( 0 );
   PyObject* self = PyObject_Call( pyclass, args, 0 );
   Py_DECREF( args );
   Py_DECREF( pyclass );

// final check before declaring success ...
   if ( ! self || ! PyROOT::ObjectProxy_Check( self ) ) {
      if ( ! PyErr_Occurred() )
         PyErr_SetString( PyExc_RuntimeError, "could not create python selector" );
      Py_XDECREF( self );
      Abort( 0 );
      return;
   }

// steal reference to new self, since the deletion will come from the C++ side
   Py_DECREF( fPySelf );
   fPySelf = self;

// inject ourselves into the base of self; destroy old identity if need be (which happens
// if the user calls the default ctor unnecessarily)
   TPySelector* oldselector = (TPySelector*)((PyROOT::ObjectProxy*)fPySelf)->fObject;
   ((PyROOT::ObjectProxy*)fPySelf)->fObject = this;
   if ( oldselector ) {
      PyROOT::TMemoryRegulator::UnregisterObject( oldselector );
      delete oldselector;
   }
}

//____________________________________________________________________________
PyObject* TPySelector::CallSelf( const char* method, PyObject* pyobject )
{
// Forward <method> to python.
   if ( ! fPySelf || fPySelf == Py_None ) {
      Py_INCREF( Py_None );
      return Py_None;
   }

   PyObject* result = 0;

// get the named method and check for python side overload by not accepting the
// binding's methodproxy
   PyObject* pymethod = PyObject_GetAttrString( fPySelf, const_cast< char* >( method ) );
   if ( ! PyROOT::MethodProxy_CheckExact( pymethod ) ) {
      if ( pyobject )
         result = PyObject_CallFunction( pymethod, const_cast< char* >( "O" ), pyobject );
      else
         result = PyObject_CallFunction( pymethod, const_cast< char* >( "" ) );
   } else {
   // silently ignore if method not overridden (note that the above can't lead
   // to a python exception, since this (TPySelector) class contains the method
   // so it is always to be found)
      Py_INCREF( Py_None );
      result = Py_None;
   }

   Py_XDECREF( pymethod );

   if ( ! result )
      Abort( 0 );

   return result;
}


//- constructors/destructor --------------------------------------------------
TPySelector::TPySelector( TTree*, PyObject* self ) : fChain( 0 ), fPySelf( 0 )
{
// Construct a TSelector derived with <self> as the underlying, which is
// generally 0 to start out with in the current PROOF framework.
   if ( self ) {
   // steal reference as this is us, as seen from python
      fPySelf = self;
   } else {
      Py_INCREF( Py_None );        // using None allows clearer diagnostics
      fPySelf = Py_None;
   }
}

//____________________________________________________________________________
TPySelector::~TPySelector()
{
// Destructor. Only deref if still holding on to Py_None (circular otherwise).
   if ( fPySelf == Py_None ) {
      Py_DECREF( fPySelf );
   }
}


//- public functions ---------------------------------------------------------
Int_t TPySelector::Version() const {
// Return version number of this selector. First forward; if not overridden, then
// yield an obvious "undefined" number, 
   PyObject* result = const_cast< TPySelector* >( this )->CallSelf( "Version" );
   if ( result && result != Py_None ) {
      Int_t ires = (Int_t)PyLong_AsLong( result );
      Py_DECREF( result );
      return ires;
   } else if ( result == Py_None ) {
      Py_DECREF( result );
   }
   return -99;
}

//____________________________________________________________________________
Int_t TPySelector::GetEntry( Long64_t entry, Int_t getall )
{
// Boilerplate get entry; same as for generated code; not forwarded.
   return fChain ? fChain->GetTree()->GetEntry( entry, getall ) : 0;
}

//____________________________________________________________________________
void TPySelector::Init( TTree* tree )
{
// Initialize with the current tree to be used; not forwarded (may be called
// multiple times, and is called from Begin() and SlaveBegin() ).
   if ( ! tree )
      return;

// set the fChain beforehand so that the python side may correct if needed
   fChain = tree;

// forward call
   PyObject* pytree = PyROOT::BindRootObject( (void*)tree, tree->IsA() );
   PyObject* result = CallSelf( "Init", pytree );
   Py_DECREF( pytree );

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );
}

//____________________________________________________________________________
Bool_t TPySelector::Notify()
{
// Forward call to derived Notify() if available.
   PyObject* result = CallSelf( "Notify" );

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );

// by default, return kTRUE, b/c the Abort will stop the processing anyway on
// a real error, so if we get here it usually means that there is no Notify()
// override on the python side of things
   return kTRUE;
}

//____________________________________________________________________________
void TPySelector::Begin( TTree* )
{
// First function called, and used to setup the python self; forward call.
   SetupPySelf();

// As per the generated code: the tree argument is deprecated (on PROOF 0 is
// passed), and hence not forwarded.
   PyObject* result = CallSelf( "Begin" );

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );
}

//____________________________________________________________________________
void TPySelector::SlaveBegin( TTree* tree )
{
// First function called on worker node, needs to make sure python self is setup,
// then store the tree to be used, initialize client, and forward call.
   SetupPySelf();
   Init( tree );

   PyObject* result = 0;
   if ( tree ) {
      PyObject* pytree = PyROOT::BindRootObject( (void*)tree, tree->IsA() );
      result = CallSelf( "SlaveBegin", pytree );
      Py_DECREF( pytree );
   } else {
      result = CallSelf( "SlaveBegin", Py_None );
   }

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );
}

//____________________________________________________________________________
Bool_t TPySelector::Process( Long64_t entry )
{
// Actual processing; call is forwarded to python self.
   if ( ! fPySelf || fPySelf == Py_None ) {
   // would like to set a python error, but can't risk that in case of a
   // configuration problem, as it would be absorbed ...

   // simply returning kFALSE will not stop processing; need to set abort
      Abort( "no python selector instance available" );
      return kFALSE;
   }

   PyObject* result = PyObject_CallMethod( fPySelf,
      const_cast< char* >( "Process" ), const_cast< char* >( "L" ), entry );
   if ( ! result ) {
      Abort( 0 );
      return kFALSE;
   }

   Bool_t bresult = (Bool_t)PyLong_AsLong( result );
   Py_DECREF( result );
   return bresult;
}

//____________________________________________________________________________
void TPySelector::SlaveTerminate()
{
// End of client; call is forwarded to python self.
   PyObject* result = CallSelf( "SlaveTerminate" );

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );
}

//____________________________________________________________________________
void TPySelector::Terminate()
{
// End of job; call is forwarded to python self.
   PyObject* result = CallSelf( "Terminate" );

   if ( ! result )
      Abort( 0 );

   Py_XDECREF( result );
}

//____________________________________________________________________________
void TPySelector::Abort( const char* why, EAbort what )
{
// If no 'why' given, read from python error
   if ( ! why && PyErr_Occurred() ) {
      PyObject *pytype = 0, *pyvalue = 0, *pytrace = 0;
      PyErr_Fetch( &pytype, &pyvalue, &pytrace );

   // abort is delayed (done at end of loop, message is current)
      PyObject* pystr = PyObject_Str( pyvalue );
      Abort( PyString_AS_STRING( pystr ), what );
      Py_DECREF( pystr );

      PyErr_Restore( pytype, pyvalue, pytrace );
   } else
      TSelector::Abort( why ? why : "", what );
}
 TPySelector.cxx:1
 TPySelector.cxx:2
 TPySelector.cxx:3
 TPySelector.cxx:4
 TPySelector.cxx:5
 TPySelector.cxx:6
 TPySelector.cxx:7
 TPySelector.cxx:8
 TPySelector.cxx:9
 TPySelector.cxx:10
 TPySelector.cxx:11
 TPySelector.cxx:12
 TPySelector.cxx:13
 TPySelector.cxx:14
 TPySelector.cxx:15
 TPySelector.cxx:16
 TPySelector.cxx:17
 TPySelector.cxx:18
 TPySelector.cxx:19
 TPySelector.cxx:20
 TPySelector.cxx:21
 TPySelector.cxx:22
 TPySelector.cxx:23
 TPySelector.cxx:24
 TPySelector.cxx:25
 TPySelector.cxx:26
 TPySelector.cxx:27
 TPySelector.cxx:28
 TPySelector.cxx:29
 TPySelector.cxx:30
 TPySelector.cxx:31
 TPySelector.cxx:32
 TPySelector.cxx:33
 TPySelector.cxx:34
 TPySelector.cxx:35
 TPySelector.cxx:36
 TPySelector.cxx:37
 TPySelector.cxx:38
 TPySelector.cxx:39
 TPySelector.cxx:40
 TPySelector.cxx:41
 TPySelector.cxx:42
 TPySelector.cxx:43
 TPySelector.cxx:44
 TPySelector.cxx:45
 TPySelector.cxx:46
 TPySelector.cxx:47
 TPySelector.cxx:48
 TPySelector.cxx:49
 TPySelector.cxx:50
 TPySelector.cxx:51
 TPySelector.cxx:52
 TPySelector.cxx:53
 TPySelector.cxx:54
 TPySelector.cxx:55
 TPySelector.cxx:56
 TPySelector.cxx:57
 TPySelector.cxx:58
 TPySelector.cxx:59
 TPySelector.cxx:60
 TPySelector.cxx:61
 TPySelector.cxx:62
 TPySelector.cxx:63
 TPySelector.cxx:64
 TPySelector.cxx:65
 TPySelector.cxx:66
 TPySelector.cxx:67
 TPySelector.cxx:68
 TPySelector.cxx:69
 TPySelector.cxx:70
 TPySelector.cxx:71
 TPySelector.cxx:72
 TPySelector.cxx:73
 TPySelector.cxx:74
 TPySelector.cxx:75
 TPySelector.cxx:76
 TPySelector.cxx:77
 TPySelector.cxx:78
 TPySelector.cxx:79
 TPySelector.cxx:80
 TPySelector.cxx:81
 TPySelector.cxx:82
 TPySelector.cxx:83
 TPySelector.cxx:84
 TPySelector.cxx:85
 TPySelector.cxx:86
 TPySelector.cxx:87
 TPySelector.cxx:88
 TPySelector.cxx:89
 TPySelector.cxx:90
 TPySelector.cxx:91
 TPySelector.cxx:92
 TPySelector.cxx:93
 TPySelector.cxx:94
 TPySelector.cxx:95
 TPySelector.cxx:96
 TPySelector.cxx:97
 TPySelector.cxx:98
 TPySelector.cxx:99
 TPySelector.cxx:100
 TPySelector.cxx:101
 TPySelector.cxx:102
 TPySelector.cxx:103
 TPySelector.cxx:104
 TPySelector.cxx:105
 TPySelector.cxx:106
 TPySelector.cxx:107
 TPySelector.cxx:108
 TPySelector.cxx:109
 TPySelector.cxx:110
 TPySelector.cxx:111
 TPySelector.cxx:112
 TPySelector.cxx:113
 TPySelector.cxx:114
 TPySelector.cxx:115
 TPySelector.cxx:116
 TPySelector.cxx:117
 TPySelector.cxx:118
 TPySelector.cxx:119
 TPySelector.cxx:120
 TPySelector.cxx:121
 TPySelector.cxx:122
 TPySelector.cxx:123
 TPySelector.cxx:124
 TPySelector.cxx:125
 TPySelector.cxx:126
 TPySelector.cxx:127
 TPySelector.cxx:128
 TPySelector.cxx:129
 TPySelector.cxx:130
 TPySelector.cxx:131
 TPySelector.cxx:132
 TPySelector.cxx:133
 TPySelector.cxx:134
 TPySelector.cxx:135
 TPySelector.cxx:136
 TPySelector.cxx:137
 TPySelector.cxx:138
 TPySelector.cxx:139
 TPySelector.cxx:140
 TPySelector.cxx:141
 TPySelector.cxx:142
 TPySelector.cxx:143
 TPySelector.cxx:144
 TPySelector.cxx:145
 TPySelector.cxx:146
 TPySelector.cxx:147
 TPySelector.cxx:148
 TPySelector.cxx:149
 TPySelector.cxx:150
 TPySelector.cxx:151
 TPySelector.cxx:152
 TPySelector.cxx:153
 TPySelector.cxx:154
 TPySelector.cxx:155
 TPySelector.cxx:156
 TPySelector.cxx:157
 TPySelector.cxx:158
 TPySelector.cxx:159
 TPySelector.cxx:160
 TPySelector.cxx:161
 TPySelector.cxx:162
 TPySelector.cxx:163
 TPySelector.cxx:164
 TPySelector.cxx:165
 TPySelector.cxx:166
 TPySelector.cxx:167
 TPySelector.cxx:168
 TPySelector.cxx:169
 TPySelector.cxx:170
 TPySelector.cxx:171
 TPySelector.cxx:172
 TPySelector.cxx:173
 TPySelector.cxx:174
 TPySelector.cxx:175
 TPySelector.cxx:176
 TPySelector.cxx:177
 TPySelector.cxx:178
 TPySelector.cxx:179
 TPySelector.cxx:180
 TPySelector.cxx:181
 TPySelector.cxx:182
 TPySelector.cxx:183
 TPySelector.cxx:184
 TPySelector.cxx:185
 TPySelector.cxx:186
 TPySelector.cxx:187
 TPySelector.cxx:188
 TPySelector.cxx:189
 TPySelector.cxx:190
 TPySelector.cxx:191
 TPySelector.cxx:192
 TPySelector.cxx:193
 TPySelector.cxx:194
 TPySelector.cxx:195
 TPySelector.cxx:196
 TPySelector.cxx:197
 TPySelector.cxx:198
 TPySelector.cxx:199
 TPySelector.cxx:200
 TPySelector.cxx:201
 TPySelector.cxx:202
 TPySelector.cxx:203
 TPySelector.cxx:204
 TPySelector.cxx:205
 TPySelector.cxx:206
 TPySelector.cxx:207
 TPySelector.cxx:208
 TPySelector.cxx:209
 TPySelector.cxx:210
 TPySelector.cxx:211
 TPySelector.cxx:212
 TPySelector.cxx:213
 TPySelector.cxx:214
 TPySelector.cxx:215
 TPySelector.cxx:216
 TPySelector.cxx:217
 TPySelector.cxx:218
 TPySelector.cxx:219
 TPySelector.cxx:220
 TPySelector.cxx:221
 TPySelector.cxx:222
 TPySelector.cxx:223
 TPySelector.cxx:224
 TPySelector.cxx:225
 TPySelector.cxx:226
 TPySelector.cxx:227
 TPySelector.cxx:228
 TPySelector.cxx:229
 TPySelector.cxx:230
 TPySelector.cxx:231
 TPySelector.cxx:232
 TPySelector.cxx:233
 TPySelector.cxx:234
 TPySelector.cxx:235
 TPySelector.cxx:236
 TPySelector.cxx:237
 TPySelector.cxx:238
 TPySelector.cxx:239
 TPySelector.cxx:240
 TPySelector.cxx:241
 TPySelector.cxx:242
 TPySelector.cxx:243
 TPySelector.cxx:244
 TPySelector.cxx:245
 TPySelector.cxx:246
 TPySelector.cxx:247
 TPySelector.cxx:248
 TPySelector.cxx:249
 TPySelector.cxx:250
 TPySelector.cxx:251
 TPySelector.cxx:252
 TPySelector.cxx:253
 TPySelector.cxx:254
 TPySelector.cxx:255
 TPySelector.cxx:256
 TPySelector.cxx:257
 TPySelector.cxx:258
 TPySelector.cxx:259
 TPySelector.cxx:260
 TPySelector.cxx:261
 TPySelector.cxx:262
 TPySelector.cxx:263
 TPySelector.cxx:264
 TPySelector.cxx:265
 TPySelector.cxx:266
 TPySelector.cxx:267
 TPySelector.cxx:268
 TPySelector.cxx:269
 TPySelector.cxx:270
 TPySelector.cxx:271
 TPySelector.cxx:272
 TPySelector.cxx:273
 TPySelector.cxx:274
 TPySelector.cxx:275
 TPySelector.cxx:276
 TPySelector.cxx:277
 TPySelector.cxx:278
 TPySelector.cxx:279
 TPySelector.cxx:280
 TPySelector.cxx:281
 TPySelector.cxx:282
 TPySelector.cxx:283
 TPySelector.cxx:284
 TPySelector.cxx:285
 TPySelector.cxx:286
 TPySelector.cxx:287
 TPySelector.cxx:288
 TPySelector.cxx:289
 TPySelector.cxx:290
 TPySelector.cxx:291
 TPySelector.cxx:292
 TPySelector.cxx:293
 TPySelector.cxx:294
 TPySelector.cxx:295
 TPySelector.cxx:296
 TPySelector.cxx:297
 TPySelector.cxx:298
 TPySelector.cxx:299
 TPySelector.cxx:300
 TPySelector.cxx:301
 TPySelector.cxx:302
 TPySelector.cxx:303
 TPySelector.cxx:304
 TPySelector.cxx:305
 TPySelector.cxx:306
 TPySelector.cxx:307
 TPySelector.cxx:308
 TPySelector.cxx:309
 TPySelector.cxx:310
 TPySelector.cxx:311
 TPySelector.cxx:312
 TPySelector.cxx:313
 TPySelector.cxx:314
 TPySelector.cxx:315
 TPySelector.cxx:316
 TPySelector.cxx:317
 TPySelector.cxx:318
 TPySelector.cxx:319
 TPySelector.cxx:320
 TPySelector.cxx:321
 TPySelector.cxx:322
 TPySelector.cxx:323
 TPySelector.cxx:324
 TPySelector.cxx:325
 TPySelector.cxx:326
 TPySelector.cxx:327
 TPySelector.cxx:328
 TPySelector.cxx:329
 TPySelector.cxx:330
 TPySelector.cxx:331
 TPySelector.cxx:332
 TPySelector.cxx:333
 TPySelector.cxx:334
 TPySelector.cxx:335
 TPySelector.cxx:336
 TPySelector.cxx:337
 TPySelector.cxx:338
 TPySelector.cxx:339
 TPySelector.cxx:340
 TPySelector.cxx:341
 TPySelector.cxx:342
 TPySelector.cxx:343
 TPySelector.cxx:344
 TPySelector.cxx:345
 TPySelector.cxx:346
 TPySelector.cxx:347
 TPySelector.cxx:348
 TPySelector.cxx:349
 TPySelector.cxx:350
 TPySelector.cxx:351
 TPySelector.cxx:352
 TPySelector.cxx:353
 TPySelector.cxx:354
 TPySelector.cxx:355
 TPySelector.cxx:356
 TPySelector.cxx:357
 TPySelector.cxx:358
 TPySelector.cxx:359
 TPySelector.cxx:360
 TPySelector.cxx:361
 TPySelector.cxx:362
 TPySelector.cxx:363
 TPySelector.cxx:364
 TPySelector.cxx:365
 TPySelector.cxx:366
 TPySelector.cxx:367
 TPySelector.cxx:368
 TPySelector.cxx:369
 TPySelector.cxx:370
 TPySelector.cxx:371
 TPySelector.cxx:372
 TPySelector.cxx:373
 TPySelector.cxx:374
 TPySelector.cxx:375