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,22 +86,40 @@ 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
-with cond:
+  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
+  import mujoco.viewer
 
-# Similar to a queue.Queue(maxsize=1), but where only one active task is allowed at a time.
+  # Similar to a queue.Queue(maxsize=1), but where only one active task is allowed at a time.
-# With queue.Queue(1), another item is allowed to be enqueued before task_done is called.
+  # With queue.Queue(1), another item is allowed to be enqueued before task_done is called.
-class _MjPythonImpl(mujoco.viewer._MjPythonBase):
+  class _MjPythonImpl(mujoco.viewer._MjPythonBase):
 
     # Termination statuses
     NOT_TERMINATED = 0
@@ -168,13 +186,13 @@ class _MjPythonImpl(mujoco.viewer._MjPythonBase):
         self._cond.notify()
 
 
-mujoco.viewer._MJPYTHON = _MjPythonImpl()
+  mujoco.viewer._MJPYTHON = _MjPythonImpl()
-atexit.register(mujoco.viewer._MJPYTHON.terminate)
+  atexit.register(mujoco.viewer._MJPYTHON.terminate)
-del _MjPythonImpl  # Don't pollute globals for user script.
 
-with cond:
+  with cond:
     cond.notify()
-del cond  # Don't pollute globals for user script.
+
+_mjpython_init()
 )", nullptr);
 
   // Run the Python interpreter main loop.
@@ -283,22 +301,29 @@ 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"(
-import ctypes
+def _mjpython_main():
+  # Don't pollute the global namespace.
+  global _mjpython_main
+  del _mjpython_main
 
-# GLFW must be initialized on the OS main thread (i.e. here).
+  import ctypes
-import glfw
-import mujoco.viewer
 
-glfw.init()
+  # GLFW must be initialized on the OS main thread (i.e. here).
-glfw.poll_events()
+  import glfw
-ctypes.CDLL(None).mjpython_hide_dock_icon()
+  import mujoco.viewer
 
-# Wait for Python main thread to finish setting up _MJPYTHON
+  glfw.init()
-with cond:
+  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:
+  while True:
     try:
       # Wait for an incoming payload.
       task = mujoco.viewer._MJPYTHON.get()
@@ -324,6 +349,8 @@ while True:
 
     finally:
       mujoco.viewer._MJPYTHON.done()
+
+_mjpython_main()
 )", nullptr);
   cpython.PyGILState_Release(gil);