From a7b43f1015a0559705f2714a680aa9f87313e603 Mon Sep 17 00:00:00 2001 From: Todd Gamblin Date: Tue, 28 Jan 2020 21:31:53 -0800 Subject: 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. --- lib/spack/spack/cmd/python.py | 16 +++++++++++++++- lib/spack/spack/test/cmd/python.py | 16 ++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) (limited to 'lib') 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 @@ -19,12 +22,23 @@ level = "long" 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 -- cgit v1.2.3-70-g09d2