Don't pollute the global namespace when using mjpython.

Fixes #2265

PiperOrigin-RevId: 713442833
Change-Id: If155d9f57d4deef8b848478ebd8077a00417d5c5

doc/changelog.rst

@@ -12,6 +12,7 @@ Python bindings
   the computation. The thread pool can be reused across calls, but then the function cannot be called simultaneously
   from multiple threads. To run multiple threaded rollouts simultaneously, use the new class ``Rollout`` which
   encapsulates the thread pool. Contribution by :github:user:`aftersomemath`.
+- Fix global namespace pollution when using ``mjpython`` (:github:issue:`2265`).
 
 General
 ^^^^^^^

python/mujoco/mjpython/mjpython.mm

@@ -86,15 +86,33 @@ void* mjpython_pymain(void* vargs) {
 
   // Set up the condition variable to pass control back to the macOS main thread.
   gil = cpython.PyGILState_Ensure();
-  cpython.PyRun_SimpleStringFlags("import threading; cond = threading.Condition()", nullptr);
+  cpython.PyRun_SimpleStringFlags(R"(
+def _mjpython_make_cond():
+  # Don't pollute the global namespace.
+  global _mjpython_make_cond
+  del _mjpython_make_cond
+
+  import threading
+
+  global cond
+  cond = threading.Condition()
+
+_mjpython_make_cond()
+)", nullptr);
   py_initialized.store(true);
 
   // Wait until GLFW is initialized on macOS main thread, set up the queue and an atexit hook
   // to enqueue a termination flag upon exit.
   cpython.PyRun_SimpleStringFlags(R"(
-import atexit
+def _mjpython_init():
+  # Don't pollute the global namespace.
+  global _mjpython_init
+  del _mjpython_init
 
-# The mujoco.viewer module should only be imported here after glfw.init() in the macOS main thread.
+  import atexit
+  import threading
+
+  # The mujoco.viewer module should only be imported after glfw.init() in the macOS main thread.
   with cond:
     cond.wait()
   import mujoco.viewer
@@ -170,11 +188,11 @@ class _MjPythonImpl(mujoco.viewer._MjPythonBase):
 
   mujoco.viewer._MJPYTHON = _MjPythonImpl()
   atexit.register(mujoco.viewer._MJPYTHON.terminate)
-del _MjPythonImpl  # Don't pollute globals for user script.
 
   with cond:
     cond.notify()
-del cond  # Don't pollute globals for user script.
+
+_mjpython_init()
 )", nullptr);
 
   // Run the Python interpreter main loop.
@@ -283,6 +301,11 @@ int main(int argc, char** argv) {
   // to finish setting up _MJPYTHON, then serve incoming viewer launch requests.
   PyGILState_STATE gil = cpython.PyGILState_Ensure();
   cpython.PyRun_SimpleStringFlags(R"(
+def _mjpython_main():
+  # Don't pollute the global namespace.
+  global _mjpython_main
+  del _mjpython_main
+
   import ctypes
 
   # GLFW must be initialized on the OS main thread (i.e. here).
@@ -294,9 +317,11 @@ glfw.poll_events()
   ctypes.CDLL(None).mjpython_hide_dock_icon()
 
   # Wait for Python main thread to finish setting up _MJPYTHON
+  global cond
   with cond:
     cond.notify()
     cond.wait()
+  del cond
 
   while True:
     try:
@@ -324,6 +349,8 @@ while True:
 
     finally:
       mujoco.viewer._MJPYTHON.done()
+
+_mjpython_main()
 )", nullptr);
   cpython.PyGILState_Release(gil);