Fix rmg_env_command dropping PYTHONPATH needed for source-tree RMG-Py (#994)

## Failure mode

`rmg_env_command()`'s `RMG_PYTHON` branch (direct-interpreter
invocation, used
on conda/mambaforge installs where micromamba's `conda` shim is broken)
unsets `_ARC_ENV_ACTIVATION_VARS` before invoking `rmg_env`'s
interpreter,
to scrub ARC's env leakage (`BABEL_LIBDIR`, `LD_LIBRARY_PATH`,
`CONDA_*`,
etc. — otherwise bound to `arc_env`'s tree and causing ABI-mismatch
crashes). `PYTHONPATH` is included in that unset list.

That's correct when RMG-Py/Arkane are pip-installed into `rmg_env`. But
on
a host where RMG-Py/Arkane are used from a **source checkout reachable
only via `PYTHONPATH`**, unsetting it and stopping there leaves the
child
interpreter unable to import `rmgpy`/`arkane` at all. Every call through
this branch dies with:

```
/home/alon/anaconda3/envs/rmg_env/bin/python: No module named arkane
```

Real-run evidence:

`/home/alon/runs/t3-pes-CH2O2/iteration_1/PDep_SA/network1_1/MSC/stderr.log`
contains exactly that line. Consequence: T3's master-equation
sensitivity
job can never run on such a host, so no P-dep network can ever qualify
for
QM refinement — the entire PDep→QM feature is silently dead, surfaced
only
as the bland "No PDep networks qualified for QM refinement".

## Why unsetting PYTHONPATH is still right

Leaving ARC's own `PYTHONPATH` bound in the child is exactly the kind of
env leakage this branch exists to scrub (see the module docstring / the
`_ARC_ENV_ACTIVATION_VARS` comment) — a stale entry (e.g. an old
checkout
on the caller's `PYTHONPATH`) could shadow `rmg_env`'s own site-packages
or ARC's own modules could bleed into the child. The unset is correct;
the
bug is stopping there instead of restoring the *target* env's own path.

## The fix

After the unset, re-export `PYTHONPATH` from ARC's own `RMG_PATH`
setting
(already resolved in `arc/settings/settings.py`, alongside
`RMG_PYTHON`),
shell-quoted like the neighbouring lines, and only when `RMG_PATH` is
truthy (never emit an empty `export PYTHONPATH=`, which would be worse
than leaving it unset):

```bash
unset CONDA_PREFIX ... PYTHONPATH PYTHONHOME
export PATH=/home/alon/anaconda3/envs/rmg_env/bin:"$PATH"
export PYTHONPATH=/home/alon/Code/RMG-Py-t3pes   # <-- restored
/home/alon/anaconda3/envs/rmg_env/bin/python -c "import arkane; print(arkane.__file__)"
# -> arkane OK
```

Verified manually on the affected host before implementing.

## The other two branches

- `MAMBA_EXE` branch and the launcher-hunt (`bash -l`) branch never
unset
`PYTHONPATH` in the first place — they route through a launcher's `run`
(which re-fires the target env's own activation hooks) or a login shell
  that sources the user's profile. Neither has this gap, so neither was
  changed.

## Test plan

- [x] Baseline: `python -m pytest arc/job/env_run_test.py` — 36 passed
(unmodified branch)
- [x] Post-fix: same command — 38 passed (36 + 2 new regression tests)
- [x] New tests assert `PYTHONPATH` is re-exported from `RMG_PATH`
(after
the `unset`) in the `RMG_PYTHON` branch, and that nothing is emitted
when `RMG_PATH` is falsy. Settings are patched via `patch.dict`, so
      the tests don't depend on this machine's real paths.
- [x] `ruff check arc/job/env_run.py arc/job/env_run_test.py` — clean

Author: calvinp0

arc/job/env_run.py

@@ -168,11 +168,16 @@ def run_in_conda_env(
 # target env, but invoking rmg_env's interpreter *directly* does not: they stay
 # pointed at arc_env's tree and make the child resolve shared-library plugins
 # against the wrong env, which is what silently kills Arkane's OpenBabel
-# import. Only the direct-interpreter branch below has to scrub them.
-# PYTHONPATH is included because the RMG helper scripts in arc/scripts/ import
-# only rmgpy/arkane/rdkit plus a sibling ``common`` module (found via the
-# script's own directory), so nothing there needs ARC on the path -- while a
-# stale entry would shadow rmg_env's site-packages.
+# import. Only the direct-interpreter branch below has to scrub this whole set;
+# the launcher branches leave the ABI-sensitive vars to ``run``'s re-activation.
+# PYTHONPATH is the exception: a launcher's ``run`` re-fires the target env's
+# activation hooks but does NOT manage PYTHONPATH, so ARC's own PYTHONPATH would
+# leak into the child on *every* branch and shadow rmg_env's site-packages.
+# ``_pythonpath_lines`` therefore normalizes it uniformly across all three
+# branches. It is kept in this tuple too so the direct-interpreter branch scrubs
+# it as part of the single ``unset``; the RMG helper scripts in arc/scripts/
+# import only rmgpy/arkane/rdkit plus a sibling ``common`` module (found via the
+# script's own directory), so nothing there needs ARC on the path.
 _ARC_ENV_ACTIVATION_VARS = (
     'CONDA_PREFIX', 'CONDA_PREFIX_1', 'CONDA_DEFAULT_ENV', 'CONDA_PROMPT_MODIFIER',
     'CONDA_SHLVL', 'LD_LIBRARY_PATH', 'BABEL_LIBDIR', 'BABEL_DATADIR',
@@ -182,6 +187,31 @@ def run_in_conda_env(
 _NO_LAUNCHER_MSG = 'micromamba, mamba, or conda is required to run RMG helper scripts'
 
 
+def _pythonpath_lines(rmg_path: str | None) -> list[str]:
+    """Bash lines that normalize ``PYTHONPATH`` for the RMG child.
+
+    A launcher's ``run`` rebinds the ABI-sensitive activation vars but never
+    manages ``PYTHONPATH``, and the direct-interpreter branch has no launcher at
+    all -- so on every branch ARC's own ``PYTHONPATH`` would otherwise leak into
+    the child. Scrub it, then re-point it at a source-tree RMG-Py/Arkane checkout
+    when ``RMG_PATH`` is configured (reachable only via ``PYTHONPATH``, rather
+    than pip-installed into rmg_env), giving all three branches one and the same
+    ``PYTHONPATH`` policy.
+
+    Args:
+        rmg_path (str | None): ARC's configured ``RMG_PATH``, or a falsy value
+                               when RMG-Py/Arkane come from rmg_env's own
+                               site-packages.
+
+    Returns:
+        list[str]: A single bash line -- ``export PYTHONPATH=<rmg_path>`` when
+                   ``rmg_path`` is truthy, else ``unset PYTHONPATH``.
+    """
+    if rmg_path:
+        return [f'export PYTHONPATH={shlex.quote(rmg_path)}']
+    return ['unset PYTHONPATH']
+
+
 def rmg_env_command(py_args: str | list[str],
                     cwd: str | None = None,
                     env_vars: dict[str, str] | None = None,
@@ -196,9 +226,19 @@ def rmg_env_command(py_args: str | list[str],
     Resolution order, preserved from the call sites this replaced:
     ``MAMBA_EXE`` (exported by setup-micromamba in CI) → ``RMG_PYTHON`` from
     ARC's settings (needed on conda/mambaforge installs where micromamba's
-    ``conda`` shim is broken) → a launcher found on PATH, hunted for under a
-    login shell so that a conda initialization block in the user's profile is
-    still honoured.
+    ``conda`` shim is broken; this branch invokes the interpreter directly, so
+    it also has to scrub ARC's leaked activation vars) → a launcher found on
+    PATH, hunted for under a login shell so that a conda initialization block in
+    the user's profile is still honoured.
+
+    Every branch normalizes ``PYTHONPATH`` identically via ``_pythonpath_lines``
+    (``RMG_PATH`` when set, scrubbed otherwise): a launcher's ``run`` re-fires
+    the target env's activation hooks but never manages ``PYTHONPATH``, so
+    without this ARC's own ``PYTHONPATH`` would leak into the child on the two
+    launcher branches -- and on all three a source-tree RMG-Py/Arkane checkout
+    reachable only via ``PYTHONPATH`` would be invisible. The direct-interpreter
+    branch additionally scrubs the ABI-sensitive vars (``LD_LIBRARY_PATH``,
+    ``BABEL_LIBDIR``, ...) that the launcher branches leave to ``run``.
 
     Args:
         py_args (str | list[str]): Everything after ``python``. Passing a
@@ -224,6 +264,7 @@ def rmg_env_command(py_args: str | list[str],
     """
     env_name = settings.get('RMG_ENV_NAME', 'rmg_env')
     rmg_python = settings.get('RMG_PYTHON')
+    rmg_path = settings.get('RMG_PATH')
     if isinstance(py_args, list):
         py_args = ' '.join(shlex.quote(arg) for arg in py_args)
 
@@ -235,23 +276,25 @@ def rmg_env_command(py_args: str | list[str],
 
     mamba_exe = os.environ.get('MAMBA_EXE', '')
     if mamba_exe and os.path.isfile(mamba_exe):
-        return '\n'.join(preamble + [
+        return '\n'.join(preamble + _pythonpath_lines(rmg_path) + [
             f'{shlex.quote(mamba_exe)} run -n {env_name} python {py_args}{suffix}',
         ])
 
     if rmg_python and os.path.isfile(rmg_python):
-        return '\n'.join(preamble + [
+        lines = preamble + [
             f'unset {" ".join(_ARC_ENV_ACTIVATION_VARS)}',
             f'export PATH={shlex.quote(os.path.dirname(rmg_python))}:"$PATH"',
-            f'{shlex.quote(rmg_python)} {py_args}{suffix}',
-        ])
+        ]
+        lines += _pythonpath_lines(rmg_path)
+        lines.append(f'{shlex.quote(rmg_python)} {py_args}{suffix}')
+        return '\n'.join(lines)
 
     # No launcher pinned by an env var and no configured interpreter: hunt for a
     # launcher on PATH. This runs under ``bash -l`` so the user's profile (where
     # conda's init block usually lives) is sourced first. The script is fed on
     # stdin via a quoted heredoc, which keeps it free of the nested shell
     # quoting the per-call-site copies of this ladder each had to get right.
-    hunted = preamble + [
+    hunted = preamble + _pythonpath_lines(rmg_path) + [
         'for _launcher in micromamba mamba conda; do',
         '    if command -v "$_launcher" >/dev/null 2>&1; then',
         f'        "$_launcher" run -n {env_name} python {py_args}{suffix}',

arc/job/env_run_test.py

@@ -384,5 +384,77 @@ def test_py_args_list_arkane_module_invocation_three_tokens(self):
         self.assertIn('python -m arkane input.py', script)
 
 
+class TestRmgEnvCommandPythonPath(unittest.TestCase):
+    """PYTHONPATH is normalized identically on all three branches: a launcher's
+    ``run`` never manages it, and the direct-interpreter branch has no launcher,
+    so ARC's own PYTHONPATH would otherwise leak into the child. When RMG_PATH is
+    set the child gets exactly that (so a source-tree RMG-Py/Arkane checkout
+    reachable only via PYTHONPATH is importable); otherwise PYTHONPATH is
+    scrubbed so rmg_env's own site-packages win."""
+
+    def setUp(self):
+        # Force the RMG_PYTHON branch: no MAMBA_EXE, and RMG_PYTHON resolves
+        # to a real file (RMG_PATH/RMG_PYTHON come from the patched settings
+        # dict below; os.path.isfile still needs a real path on disk, so
+        # point it at this test file itself).
+        self.env_patch = patch.dict(os.environ, {}, clear=False)
+        self.env_patch.start()
+        self.addCleanup(self.env_patch.stop)
+        os.environ.pop('MAMBA_EXE', None)
+        self.fake_rmg_python = __file__
+
+    def test_pythonpath_reexported_from_rmg_path(self):
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PYTHON': self.fake_rmg_python,
+                               'RMG_PATH': '/opt/RMG-Py'}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertIn(f'export PYTHONPATH={shlex.quote("/opt/RMG-Py")}', script)
+        # The re-export must come after the unset, so it is not clobbered.
+        unset_idx = script.index('unset ')
+        export_idx = script.index('export PYTHONPATH=')
+        self.assertLess(unset_idx, export_idx)
+
+    def test_no_pythonpath_export_when_rmg_path_falsy(self):
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PYTHON': self.fake_rmg_python,
+                               'RMG_PATH': None}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertNotIn('export PYTHONPATH=', script)
+
+    def test_mamba_branch_exports_pythonpath_from_rmg_path(self):
+        os.environ['MAMBA_EXE'] = __file__  # a real file forces the MAMBA_EXE branch
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PATH': '/opt/RMG-Py'}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertIn(f'export PYTHONPATH={shlex.quote("/opt/RMG-Py")}', script)
+        # Set before the launcher's ``run`` so the child inherits it.
+        self.assertLess(script.index('export PYTHONPATH='), script.index('run -n rmg_env'))
+
+    def test_mamba_branch_scrubs_pythonpath_when_rmg_path_falsy(self):
+        os.environ['MAMBA_EXE'] = __file__
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PATH': None}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertIn('unset PYTHONPATH', script)
+        self.assertNotIn('export PYTHONPATH=', script)
+
+    def test_hunt_branch_exports_pythonpath_from_rmg_path(self):
+        os.environ.pop('MAMBA_EXE', None)  # no MAMBA_EXE + no RMG_PYTHON forces the hunt branch
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PYTHON': None, 'RMG_PATH': '/opt/RMG-Py'}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertIn(f'export PYTHONPATH={shlex.quote("/opt/RMG-Py")}', script)
+        # Set inside the login-shell heredoc, before the launcher-hunt loop.
+        self.assertLess(script.index('export PYTHONPATH='), script.index('for _launcher in'))
+
+    def test_hunt_branch_scrubs_pythonpath_when_rmg_path_falsy(self):
+        os.environ.pop('MAMBA_EXE', None)
+        settings_overrides = {'RMG_ENV_NAME': 'rmg_env', 'RMG_PYTHON': None, 'RMG_PATH': None}
+        with patch.dict('arc.job.env_run.settings', settings_overrides):
+            script = rmg_env_command("-c 'pass'")
+        self.assertIn('unset PYTHONPATH', script)
+        self.assertNotIn('export PYTHONPATH=', script)
+
+
 if __name__ == '__main__':
     unittest.main()