From e3c4d827982bbaef8d61644bcd9282147c9f03ff Mon Sep 17 00:00:00 2001
From: Jason Gedge <jason.gedge@shopify.com>
Date: Fri, 28 Aug 2020 19:50:46 -0400
Subject: [PATCH] Ensure command name available for Argument completion
 locations (#2443)

---
 crates/nu-cli/src/completion/engine.rs | 18 +++++++++++++++-
 crates/nu-cli/src/shell/completer.rs   | 29 ++++++++++++--------------
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/crates/nu-cli/src/completion/engine.rs b/crates/nu-cli/src/completion/engine.rs
index 9174f828a6..e2f5cee8d1 100644
--- a/crates/nu-cli/src/completion/engine.rs
+++ b/crates/nu-cli/src/completion/engine.rs
@@ -184,6 +184,7 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
     if locations.is_empty() {
         vec![LocationType::Command.spanned(Span::unknown())]
     } else {
+        let mut command = None;
         let mut prev = None;
         for loc in locations {
             // We don't use span.contains because we want to include the end. This handles the case
@@ -210,6 +211,10 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
                 break;
             }
 
+            if let LocationType::Command = loc.item {
+                command = Some(String::from(loc.span.slice(line)));
+            }
+
             prev = Some(loc);
         }
 
@@ -221,7 +226,7 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
                 vec![LocationType::Command.spanned(Span::unknown())]
             } else {
                 // TODO this should be able to be mapped to a command
-                vec![LocationType::Argument(None, None).spanned(Span::unknown())]
+                vec![LocationType::Argument(command, None).spanned(Span::unknown())]
             }
         } else {
             // Cursor is before any possible completion location, so must be a command
@@ -354,6 +359,17 @@ mod tests {
             );
         }
 
+        #[test]
+        fn has_correct_command_name_for_argument() {
+            let registry: VecRegistry = vec![Signature::build("cd")].into();
+            let line = "cd ";
+
+            assert_eq!(
+                completion_location(line, &registry, 3),
+                vec![LocationType::Argument(Some("cd".to_string()), None)],
+            );
+        }
+
         #[test]
         fn completes_flags_with_just_a_single_hyphen() {
             let registry: VecRegistry = vec![Signature::build("du")
diff --git a/crates/nu-cli/src/shell/completer.rs b/crates/nu-cli/src/shell/completer.rs
index c45157fb02..25551efad8 100644
--- a/crates/nu-cli/src/shell/completer.rs
+++ b/crates/nu-cli/src/shell/completer.rs
@@ -45,14 +45,12 @@ impl NuCompleter {
                             flag_completer.complete(context, cmd, partial)
                         }
 
-                        LocationType::Argument(_cmd, _arg_name) => {
-                            // TODO use cmd and arg_name to narrow things down further
+                        LocationType::Argument(cmd, _arg_name) => {
                             let path_completer = crate::completion::path::Completer::new();
                             let completed_paths = path_completer.complete(context, partial);
-                            if &line[..2] == "cd" {
-                                autocomplete_only_folders(completed_paths)
-                            } else {
-                                completed_paths
+                            match cmd.as_deref().unwrap_or("") {
+                                "cd" => select_directory_suggestions(completed_paths),
+                                _ => completed_paths,
                             }
                         }
 
@@ -68,16 +66,15 @@ impl NuCompleter {
     }
 }
 
-fn autocomplete_only_folders(completed_paths: Vec<Suggestion>) -> Vec<Suggestion> {
-    let mut result = Vec::new();
-    for path in completed_paths {
-        let filepath = path.replacement.clone();
-        let md = metadata(filepath).expect("Expect filepath");
-        if md.is_dir() {
-            result.push(path);
-        }
-    }
-    result
+fn select_directory_suggestions(completed_paths: Vec<Suggestion>) -> Vec<Suggestion> {
+    completed_paths
+        .into_iter()
+        .filter(|suggestion| {
+            metadata(&suggestion.replacement)
+                .map(|md| md.is_dir())
+                .unwrap_or(false)
+        })
+        .collect()
 }
 
 fn requote(item: Suggestion) -> Suggestion {