From 335fca7049d03992aeb81e5dee64ab589013ea70 Mon Sep 17 00:00:00 2001
From: jclause <clausen.jonathan@gmail.com>
Date: Tue, 14 Jan 2025 21:29:15 -0700
Subject: [PATCH] Improve support using external modules with zsh (#48115)

* Improve support using external modules with zsh

Spack integrates with external modules by launching a python subprocess
to scrape the path from the module command. In zsh, subshells do not
inheret functions defined in the parent shell (only environment
variables). This breaks the module function in module_cmd.py.

As a workaround, source the module commands using $MODULESHOME prior to
running the module command.

* Fix formatting

* Fix flake error

* Add guard around sourcing module file

* Add improved unit testing to module src command

* Correct style

* Another attempt at style

* formatting again

* formatting again

---------

Co-authored-by: psakievich <psakiev@sandia.gov>
---
 lib/spack/spack/test/module_parsing.py | 47 ++++++++++++++++++++++++++
 lib/spack/spack/util/module_cmd.py     |  3 ++
 2 files changed, 50 insertions(+)

diff --git a/lib/spack/spack/test/module_parsing.py b/lib/spack/spack/test/module_parsing.py
index 0bbbb3c64e9..5b27742c3cd 100644
--- a/lib/spack/spack/test/module_parsing.py
+++ b/lib/spack/spack/test/module_parsing.py
@@ -35,6 +35,53 @@ def test_module_function_change_env(tmp_path):
     assert environb[b"NOT_AFFECTED"] == b"NOT_AFFECTED"
 
 
+def test_module_function_change_env_with_module_src_cmd(tmp_path):
+    environb = {
+        b"MODULESHOME": b"here",
+        b"TEST_MODULE_ENV_VAR": b"TEST_FAIL",
+        b"TEST_ANOTHER_MODULE_ENV_VAR": b"TEST_FAIL",
+        b"NOT_AFFECTED": b"NOT_AFFECTED",
+    }
+    src_file = tmp_path / "src_me"
+    src_file.write_text("export TEST_MODULE_ENV_VAR=TEST_SUCCESS\n")
+    module_src_file = tmp_path / "src_me_too"
+    module_src_file.write_text("export TEST_ANOTHER_MODULE_ENV_VAR=TEST_SUCCESS\n")
+    module("load", str(src_file), module_template=f". {src_file} 2>&1", environb=environb)
+    module(
+        "load",
+        str(src_file),
+        module_template=f". {src_file} 2>&1",
+        module_src_cmd=f". {module_src_file} 2>&1; ",
+        environb=environb,
+    )
+    assert environb[b"TEST_MODULE_ENV_VAR"] == b"TEST_SUCCESS"
+    assert environb[b"TEST_ANOTHER_MODULE_ENV_VAR"] == b"TEST_SUCCESS"
+    assert environb[b"NOT_AFFECTED"] == b"NOT_AFFECTED"
+
+
+def test_module_function_change_env_without_moduleshome_no_module_src_cmd(tmp_path):
+    environb = {
+        b"TEST_MODULE_ENV_VAR": b"TEST_FAIL",
+        b"TEST_ANOTHER_MODULE_ENV_VAR": b"TEST_FAIL",
+        b"NOT_AFFECTED": b"NOT_AFFECTED",
+    }
+    src_file = tmp_path / "src_me"
+    src_file.write_text("export TEST_MODULE_ENV_VAR=TEST_SUCCESS\n")
+    module_src_file = tmp_path / "src_me_too"
+    module_src_file.write_text("export TEST_ANOTHER_MODULE_ENV_VAR=TEST_SUCCESS\n")
+    module("load", str(src_file), module_template=f". {src_file} 2>&1", environb=environb)
+    module(
+        "load",
+        str(src_file),
+        module_template=f". {src_file} 2>&1",
+        module_src_cmd=f". {module_src_file} 2>&1; ",
+        environb=environb,
+    )
+    assert environb[b"TEST_MODULE_ENV_VAR"] == b"TEST_SUCCESS"
+    assert environb[b"TEST_ANOTHER_MODULE_ENV_VAR"] == b"TEST_FAIL"
+    assert environb[b"NOT_AFFECTED"] == b"NOT_AFFECTED"
+
+
 def test_module_function_no_change(tmpdir):
     src_file = str(tmpdir.join("src_me"))
     with open(src_file, "w", encoding="utf-8") as f:
diff --git a/lib/spack/spack/util/module_cmd.py b/lib/spack/spack/util/module_cmd.py
index ebd69f4fbee..eac754b6dfd 100644
--- a/lib/spack/spack/util/module_cmd.py
+++ b/lib/spack/spack/util/module_cmd.py
@@ -24,10 +24,13 @@
 def module(
     *args,
     module_template: Optional[str] = None,
+    module_src_cmd: Optional[str] = None,
     environb: Optional[MutableMapping[bytes, bytes]] = None,
 ):
     module_cmd = module_template or ("module " + " ".join(args))
     environb = environb or os.environb
+    if b"MODULESHOME" in environb:
+        module_cmd = module_src_cmd or "source $MODULESHOME/init/bash; " + module_cmd
 
     if args[0] in module_change_commands:
         # Suppress module output