spack python: add -m option to run modules as scripts

It's often useful to run a module with `python -m`, e.g.:

    python -m pyinstrument script.py

Running a python script this way was hard, though, as `spack python` did
not have a similar `-m` option.  This PR adds a `-m` option to `spack
python` so that we can do things like this:

    spack python -m pyinstrument ./test.py

This makes it easy to write a script that uses a small part of Spack and
then profile it.  Previously thee easiest way to do this was to write a
custom Spack command, which is often overkill.

3 files changed, 32 insertions, 2 deletions

diff --git a/lib/spack/spack/cmd/python.py b/lib/spack/spack/cmd/python.py
index 492c8f98e0..2f2290aad8 100644
--- a/lib/spack/spack/cmd/python.py
+++ b/lib/spack/spack/cmd/python.py
@@ -8,6 +8,9 @@ import sys
 import code
 import argparse
 import platform
+import runpy
+
+import llnl.util.tty as tty
 
 import spack
 
@@ -20,11 +23,22 @@ def setup_parser(subparser):
     subparser.add_argument(
         '-c', dest='python_command', help='command to execute')
     subparser.add_argument(
+        '-m', dest='module', action='store',
+        help='run library module as a script')
+    subparser.add_argument(
         'python_args', nargs=argparse.REMAINDER,
         help="file to run plus arguments")
 
 
-def python(parser, args):
+def python(parser, args, unknown_args):
+    if args.module:
+        sys.argv = ['spack-python'] + unknown_args + args.python_args
+        runpy.run_module(args.module, run_name="__main__", alter_sys=True)
+        return
+
+    if unknown_args:
+        tty.die("Unknown arguments:", " ".join(unknown_args))
+
     # Fake a main python shell by setting __name__ to __main__.
     console = code.InteractiveConsole({'__name__': '__main__',
                                        'spack': spack})
diff --git a/lib/spack/spack/test/cmd/python.py b/lib/spack/spack/test/cmd/python.py
index 074c295622..5bc05e0127 100644
--- a/lib/spack/spack/test/cmd/python.py
+++ b/lib/spack/spack/test/cmd/python.py
@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import pytest
+
 import spack
 from spack.main import SpackCommand
 
@@ -12,3 +14,17 @@ python = SpackCommand('python')
 def test_python():
     out = python('-c', 'import spack; print(spack.spack_version)')
     assert out.strip() == spack.spack_version
+
+
+def test_python_with_module():
+    # pytest rewrites a lot of modules, which interferes with runpy, so
+    # it's hard to test this.  Trying to import a module like sys, that
+    # has no code associated with it, raises an error reliably in python
+    # 2 and 3, which indicates we successfully ran runpy.run_module.
+    with pytest.raises(ImportError, match="No code object"):
+        python('-m', 'sys')
+
+
+def test_python_raises():
+    out = python('--foobar', fail_on_error=False)
+    assert "Error: Unknown arguments" in out
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index e6b7529452..b17733e1bf 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -1272,7 +1272,7 @@ _spack_pydoc() {
 _spack_python() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help -c"
+        SPACK_COMPREPLY="-h --help -c -m"
     else
         SPACK_COMPREPLY=""
     fi