Merge pull request #368 from Anatoscope/sofapython-forward-argv

[SofaPython/core] forward sys.argv to python scripts

SofaKernel/framework/sofa/helper/ArgumentParser.cpp

@@ -34,6 +34,8 @@ namespace helper
 
 typedef std::istringstream istrstream;
 
+ArgumentParser::extra_type ArgumentParser::extra;
+
 ArgumentBase::ArgumentBase(char s, string l, string h, bool m)
     : shortName(s)
     , longName(l)
@@ -86,9 +88,7 @@ ArgumentParser::~ArgumentParser()
 */
 void ArgumentParser::operator () ( int argc, char** argv )
 {
-    std::list<std::string> str;
-    for (int i=1; i<argc; ++i)
-        str.push_back(std::string(argv[i]));
+    std::list<std::string> str(argv + 1, argv + argc);
     (*this)(str);
 }
 
@@ -97,6 +97,9 @@ void ArgumentParser::operator () ( std::list<std::string> str )
     string shHelp("-");  shHelp.push_back( helpShortName );
     string lgHelp("--"); lgHelp.append( helpLongName );
     string name;
+
+    static const std::string extra_opt = "--argv";
+
     while( !str.empty() )
     {
         name = str.front();
@@ -105,16 +108,20 @@ void ArgumentParser::operator () ( std::list<std::string> str )
         // display help
         if( name == shHelp || name == lgHelp )
         {
-            if( globalHelp.size()>0 )
-                std::cout<< globalHelp <<std::endl;
-
+            if( globalHelp.size()>0 ) std::cout<< globalHelp <<std::endl;
+
             std::cout << "(short name, long name, description, default value)\n-h,\t--help: this help" << std::endl;
             std::cout << std::boolalpha;
             for( ArgVec::const_iterator a=commands.begin(), aend=commands.end(); a!=aend; ++a )
                 (*a)->print();
             std::cout << std::noboolalpha;
+
+            std::cout << "--argv [...]\t" << "forward extra args to the python interpreter" << std::endl;
+
             if( files )
                 std::cout << "others: file names" << std::endl;
+
+
             exit(EXIT_FAILURE);
         }
 
@@ -131,6 +138,12 @@ void ArgumentParser::operator () ( std::list<std::string> str )
             str.pop_front();
         }
 
+        // extra args
+        else if( name == extra_opt ) {
+            extra = extra_type(str.begin(), str.end());
+            return;
+        }
+
         // long name
         else if( name.length() > 1 && name[0]=='-' && name[1]=='-' )
         {

SofaKernel/framework/sofa/helper/ArgumentParser.h

@@ -251,13 +251,19 @@ class SOFA_HELPER_API ArgumentParser
     /// Set of remaining file
     std::vector<std::string>* files;
 
+    /// extra args appearing after --argv
+    using extra_type = std::vector<std::string>;
+    static extra_type extra;
+
     // help stuff
     string globalHelp;    ///< Overall presentation
     char helpShortName;   ///< short name for help
     string helpLongName;  ///< long name for help
 
 public:
-
+    /** last parsed extra arguments */
+    static const extra_type& extra_args() { return extra; }
+
     /// Constructor using a global help string
     ArgumentParser( const string& helpstr="", char hlpShrt='h', const string& hlpLng="help" );
 
@@ -349,6 +355,7 @@ class SOFA_HELPER_API ArgumentParser
         return (*this);
     }
 
+
     /** Parse a command line
     \param argc number of arguments + 1, as usual in C
     \param argv arguments

applications/plugins/SofaPython/PythonEnvironment.cpp

@@ -289,62 +289,65 @@ helper::logging::FileInfo::SPtr PythonEnvironment::getPythonCallingPointAsFileIn
     return SOFA_FILE_INFO_COPIED_FROM("undefined", -1);
 }
 
-bool PythonEnvironment::runFile( const char *filename, const std::vector<std::string>& arguments)
-{
-    std::string dir = sofa::helper::system::SetDirectory::GetParentDir(filename);
-    std::string bareFilename = sofa::helper::system::SetDirectory::GetFileNameWithoutExtension(filename);
-
-    if(!arguments.empty())
-    {
-        char**argv = new char*[arguments.size()+1];
-        argv[0] = new char[bareFilename.size()+1];
-        strcpy( argv[0], bareFilename.c_str() );
-        for( size_t i=0 ; i<arguments.size() ; ++i )
-        {
-            argv[i+1] = new char[arguments[i].size()+1];
-            strcpy( argv[i+1], arguments[i].c_str() );
+bool PythonEnvironment::runFile( const char *filename, const std::vector<std::string>& arguments) {
+    const std::string dir = sofa::helper::system::SetDirectory::GetParentDir(filename);
+
+    // pro-tip: FileNameWithoutExtension == basename
+    const std::string basename = sofa::helper::system::SetDirectory::GetFileNameWithoutExtension(filename);
+
+    // setup sys.argv if needed
+    if(!arguments.empty() ) {
+        std::vector<const char*> argv;
+        argv.push_back(basename.c_str());
+
+        for(const std::string& arg : arguments) {
+            argv.push_back(arg.c_str());
         }
-
-        Py_SetProgramName(argv[0]); // TODO check what it is doing exactly
-        PySys_SetArgv(arguments.size()+1, argv);
-
-        for( size_t i=0 ; i<arguments.size()+1 ; ++i )
-        {
-            delete [] argv[i];
-        }
-        delete [] argv;
+
+        Py_SetProgramName((char*) argv[0]); // TODO check what it is doing exactly
+        PySys_SetArgv(argv.size(), (char**)argv.data());
     }
-
+    
     // Load the scene script
-    char* pythonFilename = strdup(filename);
-    PyObject* scriptPyFile = PyFile_FromString(pythonFilename, (char*)("r"));
-    free(pythonFilename);
-
-    if( !scriptPyFile )
-    {
+    PyObject* script = PyFile_FromString((char*)filename, (char*)("r"));
+
+    if( !script ) {
         SP_MESSAGE_ERROR("cannot open file:" << filename)
         PyErr_Print();
         return false;
     }
-
-    PyObject* pDict = PyModule_GetDict(PyImport_AddModule("__main__"));
-
-    std::string backupFileName;
-    PyObject* backupFileObject = PyDict_GetItemString(pDict, "__file__");
-    if(backupFileObject)
-        backupFileName = PyString_AsString(backupFileObject);
-
-    PyObject* newFileObject = PyString_FromString(filename);
-    PyDict_SetItemString(pDict, "__file__", newFileObject);
-
-    int error = PyRun_SimpleFileEx(PyFile_AsFile(scriptPyFile), filename, 1);
-
-    backupFileObject = PyString_FromString(backupFileName.c_str());
-    PyDict_SetItemString(pDict, "__file__", backupFileObject);
-
-    if(0 != error)
+
+    PyObject* __main__ = PyModule_GetDict(PyImport_AddModule("__main__"));
+
+    // save/restore __main__.__file__
+    PyObject* __file__ = PyDict_GetItemString(__main__, "__file__");
+    Py_XINCREF(__file__);
+
+    // temporarily set __main__.__file__ = filename during file loading
     {
-        SP_MESSAGE_ERROR("Script (file:" << bareFilename << ") import error")
+        PyObject* __tmpfile__ = PyString_FromString(filename);
+        PyDict_SetItemString(__main__, "__file__", __tmpfile__);
+        Py_XDECREF(__tmpfile__);
+    }
+
+    const int error = PyRun_SimpleFileEx(PyFile_AsFile(script), filename, 0);
+
+    // don't wait for gc to close the file
+    PyObject_CallMethod(script, (char*) "close", NULL);
+    Py_XDECREF(script);
+
+    // restore backup if needed
+    if(__file__) {
+        PyDict_SetItemString(__main__, "__file__", __file__);
+    } else {
+        const int err = PyDict_DelItemString(__main__, "__file__");
+        assert(!err); (void) err;
+    }
+
+    Py_XDECREF(__file__);  
+
+    if(error) {
+        SP_MESSAGE_ERROR("Script (file:" << basename << ") import error")
         PyErr_Print();
         return false;
     }

applications/plugins/SofaPython/SceneLoaderPY.cpp

@@ -25,6 +25,7 @@
 
 
 #include <sofa/simulation/Simulation.h>
+#include <sofa/helper/ArgumentParser.h>
 #include <SofaSimulationCommon/xml/NodeElement.h>
 #include <SofaSimulationCommon/FindByTypeVisitor.h>
 
@@ -83,7 +84,7 @@ void SceneLoaderPY::getExtensionList(ExtensionList* list)
 sofa::simulation::Node::SPtr SceneLoaderPY::load(const char *filename)
 {
     sofa::simulation::Node::SPtr root;
-    loadSceneWithArguments(filename, {}, &root);
+    loadSceneWithArguments(filename, helper::ArgumentParser::extra_args(), &root);
     return root;
 }