spack python: allow use of IPython (#20329)

This adds a -i option to "spack python" which allows use of the
IPython interpreter; it can be used with "spack python -i ipython".
This assumes it is available in the Python instance used to run
Spack (i.e. that you can "import IPython").

2 files changed, 76 insertions, 12 deletions

diff --git a/lib/spack/docs/developer_guide.rst b/lib/spack/docs/developer_guide.rst
index 6df47cec87..9964bb6e51 100644
--- a/lib/spack/docs/developer_guide.rst
+++ b/lib/spack/docs/developer_guide.rst
@@ -396,20 +396,42 @@ other Spack modules:
    True
    >>>
 
-You can also run a single command:
+If you prefer using an IPython interpreter, given that IPython is installed
+you can specify the interpreter with ``-i``:
+
+.. code-block:: console
+
+   $ spack python -i ipython
+   Python 3.8.3 (default, May 19 2020, 18:47:26) 
+   Type 'copyright', 'credits' or 'license' for more information
+   IPython 7.17.0 -- An enhanced Interactive Python. Type '?' for help.
+
+
+   Spack version 0.16.0
+   Python 3.8.3, Linux x86_64
+
+   In [1]:
+
+
+With either interpreter you can run a single command:
 
 .. code-block:: console
 
    $ spack python -c 'import distro; distro.linux_distribution()'
-   ('Fedora', '25', 'Workstation Edition')
+   ('Ubuntu', '18.04', 'Bionic Beaver')
+
+   $ spack python -i ipython -c 'import distro; distro.linux_distribution()'
+   Out[1]: ('Ubuntu', '18.04', 'Bionic Beaver')
 
 or a file:
 
 .. code-block:: console
 
    $ spack python ~/test_fetching.py
+   $ spack python -i ipython ~/test_fetching.py
+
+just like you would with the normal ``python`` command. 
 
-just like you would with the normal ``python`` command.
 
 .. _cmd-spack-url:
 
diff --git a/lib/spack/spack/cmd/python.py b/lib/spack/spack/cmd/python.py
index dac2d15b43..588ddce853 100644
--- a/lib/spack/spack/cmd/python.py
+++ b/lib/spack/spack/cmd/python.py
@@ -28,6 +28,9 @@ def setup_parser(subparser):
     subparser.add_argument(
         '-c', dest='python_command', help='command to execute')
     subparser.add_argument(
+        '-i', dest='python_interpreter', help='python interpreter',
+        choices=['python', 'ipython'], default='python')
+    subparser.add_argument(
         '-m', dest='module', action='store',
         help='run library module as a script')
     subparser.add_argument(
@@ -48,24 +51,63 @@ def python(parser, args, unknown_args):
     if unknown_args:
         tty.die("Unknown arguments:", " ".join(unknown_args))
 
+    # Unexpected behavior from supplying both
+    if args.python_command and args.python_args:
+        tty.die("You can only specify a command OR script, but not both.")
+
+    # Run user choice of interpreter
+    if args.python_interpreter == "ipython":
+        return spack.cmd.python.ipython_interpreter(args)
+    return spack.cmd.python.python_interpreter(args)
+
+
+def ipython_interpreter(args):
+    """An ipython interpreter is intended to be interactive, so it doesn't
+    support running a script or arguments
+    """
+    try:
+        import IPython
+    except ImportError:
+        tty.die("ipython is not installed, install and try again.")
+
+    if "PYTHONSTARTUP" in os.environ:
+        startup_file = os.environ["PYTHONSTARTUP"]
+        if os.path.isfile(startup_file):
+            with open(startup_file) as startup:
+                exec(startup.read())
+
+    # IPython can also support running a script OR command, not both
+    if args.python_args:
+        IPython.start_ipython(argv=args.python_args)
+    elif args.python_command:
+        IPython.start_ipython(argv=['-c', args.python_command])
+    else:
+        header = ("Spack version %s\nPython %s, %s %s"
+                  % (spack.spack_version, platform.python_version(),
+                     platform.system(), platform.machine()))
+
+        __name__ = "__main__"  # noqa
+        IPython.embed(module="__main__", header=header)
+
+
+def python_interpreter(args):
+    """A python interpreter is the default interpreter
+    """
     # Fake a main python shell by setting __name__ to __main__.
     console = code.InteractiveConsole({'__name__': '__main__',
                                        'spack': spack})
-
     if "PYTHONSTARTUP" in os.environ:
         startup_file = os.environ["PYTHONSTARTUP"]
         if os.path.isfile(startup_file):
             with open(startup_file) as startup:
                 console.runsource(startup.read(), startup_file, 'exec')
 
-    python_args = args.python_args
-    python_command = args.python_command
-    if python_command:
-        console.runsource(python_command)
-    elif python_args:
-        sys.argv = python_args
-        with open(python_args[0]) as file:
-            console.runsource(file.read(), python_args[0], 'exec')
+    if args.python_command:
+        console.runsource(args.python_command)
+    elif args.python_args:
+        sys.argv = args.python_args
+        with open(args.python_args[0]) as file:
+            console.runsource(file.read(), args.python_args[0], 'exec')
     else:
         # Provides readline support, allowing user to use arrow keys
         console.push('import readline')