studio: share ImGui and ImPlot contexts across python child extension modules.

Each pybind extension module holds its own copy of the ImGui/ImPlot globals (extension modules are loaded RTLD_LOCAL), so the viewer that owns the contexts must share the pointers explicitly. Without this, client plotting GUIs and user plotting scripts (e.g. the implot sample) crash on null context pointers.

PiperOrigin-RevId: 951917287
Change-Id: I639a173c2ae9f1af4feaf57aff7b5c81854b7393

python/mujoco/experimental/implot/implot.cc

@@ -33,6 +33,20 @@ PYBIND11_MODULE(implot, m) {
   // Import dear_imgui to make types like ImVec2 available.
   py::module_::import("mujoco.experimental.dear_imgui.dear_imgui");
 
+  // Each pybind module holds its own ImGui/ImPlot context globals
+  m.def(
+      "set_imgui_context",
+      [](intptr_t ptr) {
+        ImGui::SetCurrentContext(reinterpret_cast<ImGuiContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImGui context pointer.");
+  m.def(
+      "set_implot_context",
+      [](intptr_t ptr) {
+        ImPlot::SetCurrentContext(reinterpret_cast<ImPlotContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImPlot context pointer.");
+
   // Types.
   py::class_<ImPlotPoint>(m, "Point")
       .def(py::init<>())