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

@@ -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:

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