Add "spack [cd|location] --source-dir" (#22321)

2 files changed, 80 insertions, 56 deletions

diff --git a/lib/spack/spack/cmd/location.py b/lib/spack/spack/cmd/location.py
index 24236d7e83..2461074d88 100644
--- a/lib/spack/spack/cmd/location.py
+++ b/lib/spack/spack/cmd/location.py
@@ -48,8 +48,12 @@ def setup_parser(subparser):
         '-S', '--stages', action='store_true',
         help="top level stage directory")
     directories.add_argument(
+        '--source-dir', action='store_true',
+        help="source directory for a spec "
+             "(requires it to be staged first)")
+    directories.add_argument(
         '-b', '--build-dir', action='store_true',
-        help="checked out or expanded source directory for a spec "
+        help="build directory for a spec "
              "(requires it to be staged first)")
     directories.add_argument(
         '-e', '--env', action='store',
@@ -61,64 +65,78 @@ def setup_parser(subparser):
 def location(parser, args):
     if args.module_dir:
         print(spack.paths.module_path)
+        return
 
-    elif args.spack_root:
+    if args.spack_root:
         print(spack.paths.prefix)
+        return
 
-    elif args.env:
+    if args.env:
         path = spack.environment.root(args.env)
         if not os.path.isdir(path):
             tty.die("no such environment: '%s'" % args.env)
         print(path)
+        return
 
-    elif args.packages:
+    if args.packages:
         print(spack.repo.path.first_repo().root)
+        return
 
-    elif args.stages:
+    if args.stages:
         print(spack.stage.get_stage_root())
-
-    else:
-        specs = spack.cmd.parse_specs(args.spec)
-        if not specs:
-            tty.die("You must supply a spec.")
-        if len(specs) != 1:
-            tty.die("Too many specs.  Supply only one.")
-
-        if args.install_dir:
-            # install_dir command matches against installed specs.
-            env = ev.get_env(args, 'location')
-            spec = spack.cmd.disambiguate_spec(specs[0], env)
-            print(spec.prefix)
-
-        else:
-            spec = specs[0]
-
-            if args.package_dir:
-                # This one just needs the spec name.
-                print(spack.repo.path.dirname_for_package_name(spec.name))
-
-            else:
-                spec = spack.cmd.matching_spec_from_env(spec)
-                pkg = spec.package
-
-                if args.stage_dir:
-                    print(pkg.stage.path)
-
-                else:  # args.build_dir is the default.
-                    if not pkg.stage.expanded:
-                        tty.die("Build directory does not exist yet. "
-                                "Run this to create it:",
-                                "spack stage " + " ".join(args.spec))
-
-                    # Out of source builds have build_directory defined
-                    if hasattr(pkg, 'build_directory'):
-                        # build_directory can be either absolute or relative
-                        # to the stage path in either case os.path.join makes it
-                        # absolute
-                        print(os.path.normpath(os.path.join(
-                            pkg.stage.path,
-                            pkg.build_directory
-                        )))
-                    else:
-                        # Otherwise assume in-source builds
-                        return print(pkg.stage.source_path)
+        return
+
+    specs = spack.cmd.parse_specs(args.spec)
+
+    if not specs:
+        tty.die("You must supply a spec.")
+
+    if len(specs) != 1:
+        tty.die("Too many specs.  Supply only one.")
+
+    # install_dir command matches against installed specs.
+    if args.install_dir:
+        env = ev.get_env(args, 'location')
+        spec = spack.cmd.disambiguate_spec(specs[0], env)
+        print(spec.prefix)
+        return
+
+    spec = specs[0]
+
+    # Package dir just needs the spec name
+    if args.package_dir:
+        print(spack.repo.path.dirname_for_package_name(spec.name))
+        return
+
+    # Either concretize or filter from already concretized environment
+    spec = spack.cmd.matching_spec_from_env(spec)
+    pkg = spec.package
+
+    if args.stage_dir:
+        print(pkg.stage.path)
+        return
+
+    if args.build_dir:
+        # Out of source builds have build_directory defined
+        if hasattr(pkg, 'build_directory'):
+            # build_directory can be either absolute or relative to the stage path
+            # in either case os.path.join makes it absolute
+            print(os.path.normpath(os.path.join(
+                pkg.stage.path,
+                pkg.build_directory
+            )))
+            return
+
+        # Otherwise assume in-source builds
+        print(pkg.stage.source_path)
+        return
+
+    # source and build dir remain, they require the spec to be staged
+    if not pkg.stage.expanded:
+        tty.die("Source directory does not exist yet. "
+                "Run this to create it:",
+                "spack stage " + " ".join(args.spec))
+
+    if args.source_dir:
+        print(pkg.stage.source_path)
+        return
diff --git a/lib/spack/spack/test/cmd/location.py b/lib/spack/spack/test/cmd/location.py
index 4141c4fbbe..72c3912503 100644
--- a/lib/spack/spack/test/cmd/location.py
+++ b/lib/spack/spack/test/cmd/location.py
@@ -52,18 +52,24 @@ def test_location_build_dir(mock_spec):
     assert location('--build-dir', spec.name).strip() == pkg.stage.source_path
 
 
-def test_location_build_dir_missing():
-    """Tests spack location --build-dir with a missing build directory."""
+def test_location_source_dir(mock_spec):
+    """Tests spack location --source-dir."""
+    spec, pkg = mock_spec
+    assert location('--source-dir', spec.name).strip() == pkg.stage.source_path
+
+
+def test_location_source_dir_missing():
+    """Tests spack location --source-dir with a missing source directory."""
     spec = 'mpileaks'
     prefix = "==> Error: "
-    expected = "%sBuild directory does not exist yet. Run this to create it:"\
+    expected = "%sSource directory does not exist yet. Run this to create it:"\
                "%s  spack stage %s" % (prefix, os.linesep, spec)
-    out = location('--build-dir', spec, fail_on_error=False).strip()
+    out = location('--source-dir', spec, fail_on_error=False).strip()
     assert out == expected
 
 
 @pytest.mark.parametrize('options', [([]),
-                                     (['--build-dir', 'mpileaks']),
+                                     (['--source-dir', 'mpileaks']),
                                      (['--env', 'missing-env']),
                                      (['spec1', 'spec2'])])
 def test_location_cmd_error(options):