wb_robot_cleanup fix for python plugins

webots_ros2_driver/include/webots_ros2_driver/PythonPlugin.hpp

@@ -27,6 +27,7 @@ namespace webots_ros2_driver {
   public:
     void init(webots_ros2_driver::WebotsNode *node, std::unordered_map<std::string, std::string> &parameters) override;
     void step() override;
+    void stop();
 
     static std::shared_ptr<PythonPlugin> createFromType(const std::string &type);
 

webots_ros2_driver/include/webots_ros2_driver/WebotsNode.hpp

@@ -40,6 +40,7 @@ namespace webots_ros2_driver {
     int step();
     std::string urdf() const { return mRobotDescription; };
     static void handleSignals();
+    std::vector<std::shared_ptr<PluginInterface>> mPlugins;
 
   private:
     std::unordered_map<std::string, std::string> getDeviceRosProperties(const std::string &name) const;
@@ -57,7 +58,6 @@ namespace webots_ros2_driver {
 
     rclcpp::TimerBase::SharedPtr mTimer;
     int mStep;
-    std::vector<std::shared_ptr<PluginInterface>> mPlugins;
     pluginlib::ClassLoader<PluginInterface> mPluginLoader;
     tinyxml2::XMLElement *mWebotsXMLElement;
     std::shared_ptr<tinyxml2::XMLDocument> mRobotDescriptionDocument;

webots_ros2_driver/src/Driver.cpp

@@ -16,6 +16,8 @@
 #include <webots/vehicle/driver.h>
 #include <memory>
 #include <rclcpp/rclcpp.hpp>
+#include <webots_ros2_driver/PluginInterface.hpp>
+#include <webots_ros2_driver/PythonPlugin.hpp>
 #include <webots_ros2_driver/WebotsNode.hpp>
 
 int main(int argc, char **argv) {
@@ -44,7 +46,21 @@ int main(int argc, char **argv) {
       break;
     rclcpp::spin_some(node);
   }
-  wb_robot_cleanup();
+  // Check if the plugin is actually an instance of PythonPlugin
+  bool isPython = false;
+  for (std::shared_ptr<webots_ros2_driver::PluginInterface> plugin : node->mPlugins) {
+    std::shared_ptr<webots_ros2_driver::PythonPlugin> pythonPlugin =
+      std::dynamic_pointer_cast<webots_ros2_driver::PythonPlugin>(plugin);
+    if (pythonPlugin) {
+      pythonPlugin->stop();  // stop only for python plugins
+      isPython = true;
+      break;
+    }
+  }
+  if (!isPython) {
+    wb_robot_cleanup();
+  }
+  RCLCPP_INFO(node->get_logger(), "Controller successfully disconnected from robot in Webots simulation.");
   rclcpp::shutdown();
   return 0;
 }

webots_ros2_driver/src/PythonPlugin.cpp

@@ -20,6 +20,11 @@ namespace webots_ros2_driver {
     PyErr_Print();
   }
 
+  void PythonPlugin::stop() {
+    PyObject_CallMethod(mPyPlugin, "stop", "");
+    Py_Finalize();
+  }
+
   PyObject *PythonPlugin::getPyWebotsNodeInstance() {
     if (gPyWebotsNode)
       return gPyWebotsNode;
@@ -37,6 +42,7 @@ from controller import Supervisor
 sys.path.insert(1, os.path.dirname(controller.__file__))
 from vehicle import Driver
 
+
 class WebotsNode:
     def __init__(self):
         if Driver.isInitialisationPossible():

webots_ros2_driver/src/WebotsNode.cpp

@@ -32,8 +32,6 @@
 #include <webots_ros2_driver/plugins/static/Ros2RangeFinder.hpp>
 #include <webots_ros2_driver/plugins/static/Ros2Receiver.hpp>
 #include <webots_ros2_driver/plugins/static/Ros2VacuumGripper.hpp>
-#include "webots_ros2_driver/PluginInterface.hpp"
-
 #include "webots_ros2_driver/PluginInterface.hpp"
 #include "webots_ros2_driver/PythonPlugin.hpp"
 
@@ -258,6 +256,7 @@ namespace webots_ros2_driver {
     if (gShutdownSignalReceived && !mWaitingForUrdfRobotToBeRemoved) {
       mRemoveUrdfRobotPublisher->publish(mRemoveUrdfRobotMessage);
       mWaitingForUrdfRobotToBeRemoved = true;
+      return -1;
     }
 
     const int result = wb_robot_step(mStep);