From 496f6b47fcdbae11259716b822f914044e10763a Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Fri, 13 Feb 2026 03:43:58 -0800
Subject: [PATCH] In mjpython, check the result of Py_InitializeFromConfig.

This is to properly handle the case where the interpreter wants to exit immediately, e.g. when `mjpython --version` is run.

Currently, we call `PyGILState_Ensure` regardless of the result of `Py_InitializeFromConfig` resulting in a segfault, as in issue #2870.

PiperOrigin-RevId: 869652555
Change-Id: I3a898279eb8aa25afbaece904cfdcc4fdcc92972
---
 python/mujoco/mjpython/mjpython.mm | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/python/mujoco/mjpython/mjpython.mm b/python/mujoco/mjpython/mjpython.mm
index f7e159c1..47979b8b 100644
--- a/python/mujoco/mjpython/mjpython.mm
+++ b/python/mujoco/mjpython/mjpython.mm
@@ -57,6 +57,8 @@ struct {
   CPYTHON_FN(PyGILState_Ensure);
   CPYTHON_FN(PyGILState_Release);
   CPYTHON_FN(PyRun_SimpleStringFlags);
+  CPYTHON_FN(PyStatus_Exception);
+  CPYTHON_FN(Py_ExitStatusException);
   CPYTHON_FN(Py_FinalizeEx);
   CPYTHON_FN(Py_InitializeFromConfig);
   CPYTHON_FN(Py_RunMain);
@@ -65,7 +67,7 @@ struct {
 #undef CPYTHON_FN
 } cpython;
 
-std::atomic_bool py_initialized = false;
+std::atomic_bool py_initialized;
 
 struct Args {
   int argc;
@@ -80,8 +82,20 @@ void* mjpython_pymain(void* vargs) {
   // Initialize the Python interpreter.
   PyConfig config;
   cpython.PyConfig_InitPythonConfig(&config);
-  cpython.PyConfig_SetBytesArgv(&config, args->argc, args->argv);
-  cpython.Py_InitializeFromConfig(&config);
+  {
+    PyStatus status = cpython.PyConfig_SetBytesArgv(&config, args->argc, args->argv);
+    if (cpython.PyStatus_Exception(status)) {
+      // Calls exit() which aborts the entire process.
+      cpython.Py_ExitStatusException(status);
+    }
+  }
+  {
+    PyStatus status = cpython.Py_InitializeFromConfig(&config);
+    if (cpython.PyStatus_Exception(status)) {
+      // Calls exit() which aborts the entire process.
+      cpython.Py_ExitStatusException(status);
+    }
+  }
   cpython.PyConfig_Clear(&config);
 
   // Set up the condition variable to pass control back to the macOS main thread.
@@ -205,6 +219,7 @@ _mjpython_init()
 }  // namespace
 
 int main(int argc, char** argv) {
+  py_initialized.store(false);
   const char* libpython_path = getenv("MJPYTHON_LIBPYTHON");
   if (!libpython_path || !libpython_path[0]) {
     std::cerr << "This binary must be launched via the mjpython.py script.\n";