nu-cli/completions: completion for use and source (#5210)

* nu-cli/completions: completion for use and source

* handle subfolders for different base dirs

* fix clippy errors

crates/nu-cli/src/completions/completer.rs

@@ -1,6 +1,6 @@
 use crate::completions::{
-    CommandCompletion, Completer, CustomCompletion, FileCompletion, FlagCompletion,
-    VariableCompletion,
+    CommandCompletion, Completer, CustomCompletion, DotNuCompletion, FileCompletion,
+    FlagCompletion, VariableCompletion,
 };
 use nu_parser::{flatten_expression, parse, FlatShape};
 use nu_protocol::{
@@ -84,6 +84,30 @@ impl NuCompleter {
                         let mut prefix = working_set.get_span_contents(flat.0).to_vec();
                         prefix.remove(pos - flat.0.start);
 
+                        // Completions that depends on the previous expression (e.g: use, source)
+                        if flat_idx > 0 {
+                            if let Some(previous_expr) = flattened.get(flat_idx - 1) {
+                                // Read the content for the previous expression
+                                let prev_expr_str =
+                                    working_set.get_span_contents(previous_expr.0).to_vec();
+
+                                // Completion for .nu files
+                                if prev_expr_str == b"use" || prev_expr_str == b"source" {
+                                    let mut completer =
+                                        DotNuCompletion::new(self.engine_state.clone());
+
+                                    return self.process_completion(
+                                        &mut completer,
+                                        &working_set,
+                                        prefix,
+                                        new_span,
+                                        offset,
+                                        pos,
+                                    );
+                                }
+                            }
+                        }
+
                         // Variables completion
                         if prefix.starts_with(b"$") || most_left_var.is_some() {
                             let mut completer = VariableCompletion::new(

crates/nu-cli/src/completions/dotnu_completions.rs

@@ -0,0 +1,126 @@
+use crate::completions::{
+    file_path_completion, partial_from, Completer, CompletionOptions, SortBy,
+};
+use nu_protocol::{
+    engine::{EngineState, StateWorkingSet},
+    Span,
+};
+use reedline::Suggestion;
+use std::sync::Arc;
+const SEP: char = std::path::MAIN_SEPARATOR;
+
+#[derive(Clone)]
+pub struct DotNuCompletion {
+    engine_state: Arc<EngineState>,
+}
+
+impl DotNuCompletion {
+    pub fn new(engine_state: Arc<EngineState>) -> Self {
+        Self { engine_state }
+    }
+}
+
+impl Completer for DotNuCompletion {
+    // Replace base filter with no filter once all the results are already filtered
+    fn filter(&self, _: Vec<u8>, items: Vec<Suggestion>, _: CompletionOptions) -> Vec<Suggestion> {
+        items
+    }
+
+    fn fetch(
+        &mut self,
+        _: &StateWorkingSet,
+        prefix: Vec<u8>,
+        span: Span,
+        offset: usize,
+        _: usize,
+    ) -> (Vec<Suggestion>, CompletionOptions) {
+        let prefix_str = String::from_utf8_lossy(&prefix).to_string();
+        let mut search_dirs: Vec<String> = vec![];
+        let (base_dir, mut partial) = partial_from(&prefix_str);
+        let mut is_current_folder = false;
+
+        // Fetch the lib dirs
+        let lib_dirs: Vec<String> =
+            if let Some(lib_dirs) = self.engine_state.env_vars.get("NU_LIB_DIRS") {
+                lib_dirs
+                    .as_list()
+                    .into_iter()
+                    .flat_map(|it| {
+                        it.iter().map(|x| {
+                            x.as_path()
+                                .expect("internal error: failed to convert lib path")
+                        })
+                    })
+                    .map(|it| {
+                        it.into_os_string()
+                            .into_string()
+                            .expect("internal error: failed to convert OS path")
+                    })
+                    .collect()
+            } else {
+                vec![]
+            };
+
+        // Check if the base_dir is a folder
+        if base_dir != "./" {
+            // Add the base dir into the directories to be searched
+            search_dirs.push(base_dir.clone());
+
+            // Reset the partial adding the basic dir back
+            // in order to make the span replace work properly
+            let mut base_dir_partial = base_dir;
+            base_dir_partial.push_str(&partial);
+
+            partial = base_dir_partial;
+        } else {
+            // Fetch the current folder
+            let current_folder = if let Some(d) = self.engine_state.env_vars.get("PWD") {
+                match d.as_string() {
+                    Ok(s) => s,
+                    Err(_) => "".to_string(),
+                }
+            } else {
+                "".to_string()
+            };
+            is_current_folder = true;
+
+            // Add the current folder and the lib dirs into the
+            // directories to be searched
+            search_dirs.push(current_folder);
+            search_dirs.extend(lib_dirs);
+        }
+
+        // Fetch the files filtering the ones that ends with .nu
+        // and transform them into suggestions
+        let output: Vec<Suggestion> = search_dirs
+            .into_iter()
+            .flat_map(|it| {
+                file_path_completion(span, &partial, &it)
+                    .into_iter()
+                    .filter(|it| {
+                        // Different base dir, so we list the .nu files or folders
+                        if !is_current_folder {
+                            it.1.ends_with(".nu") || it.1.ends_with(SEP)
+                        } else {
+                            // Lib dirs, so we filter only the .nu files
+                            it.1.ends_with(".nu")
+                        }
+                    })
+                    .map(move |x| Suggestion {
+                        value: x.1,
+                        description: None,
+                        extra: None,
+                        span: reedline::Span {
+                            start: x.0.start - offset,
+                            end: x.0.end - offset,
+                        },
+                    })
+            })
+            .collect();
+
+        // Options
+        let options = CompletionOptions::new(false, true, SortBy::LevenshteinDistance);
+
+        (output, options)
+    }
+}

crates/nu-cli/src/completions/file_completions.rs

@@ -104,23 +104,26 @@ impl Completer for FileCompletion {
     }
 }
 
+pub fn partial_from(input: &str) -> (String, String) {
+    let partial = input.replace('\'', "");
+
+    // If partial is only a word we want to search in the current dir
+    let (base, rest) = partial.rsplit_once(is_separator).unwrap_or((".", &partial));
+    // On windows, this standardizes paths to use \
+    let mut base = base.replace(is_separator, &SEP.to_string());
+
+    // rsplit_once removes the separator
+    base.push(SEP);
+
+    (base.to_string(), rest.to_string())
+}
+
 pub fn file_path_completion(
     span: nu_protocol::Span,
     partial: &str,
     cwd: &str,
 ) -> Vec<(nu_protocol::Span, String)> {
-    let partial = partial.replace('\'', "");
-
-    let (base_dir_name, partial) = {
-        // If partial is only a word we want to search in the current dir
-        let (base, rest) = partial.rsplit_once(is_separator).unwrap_or((".", &partial));
-        // On windows, this standardizes paths to use \
-        let mut base = base.replace(is_separator, &SEP.to_string());
-
-        // rsplit_once removes the separator
-        base.push(SEP);
-        (base, rest)
-    };
+    let (base_dir_name, partial) = partial_from(partial);
 
     let base_dir = nu_path::expand_path_with(&base_dir_name, cwd);
     // This check is here as base_dir.read_dir() with base_dir == "" will open the current dir
@@ -134,7 +137,7 @@ pub fn file_path_completion(
             .filter_map(|entry| {
                 entry.ok().and_then(|entry| {
                     let mut file_name = entry.file_name().to_string_lossy().into_owned();
-                    if matches(partial, &file_name) {
+                    if matches(&partial, &file_name) {
                         let mut path = format!("{}{}", base_dir_name, file_name);
                         if entry.path().is_dir() {
                             path.push(SEP);

crates/nu-cli/src/completions/mod.rs

@@ -3,6 +3,7 @@ mod command_completions;
 mod completer;
 mod completion_options;
 mod custom_completions;
+mod dotnu_completions;
 mod file_completions;
 mod flag_completions;
 mod variable_completions;
@@ -12,6 +13,7 @@ pub use command_completions::CommandCompletion;
 pub use completer::NuCompleter;
 pub use completion_options::{CompletionOptions, SortBy};
 pub use custom_completions::CustomCompletion;
-pub use file_completions::{file_path_completion, FileCompletion};
+pub use dotnu_completions::DotNuCompletion;
+pub use file_completions::{file_path_completion, partial_from, FileCompletion};
 pub use flag_completions::FlagCompletion;
 pub use variable_completions::VariableCompletion;

crates/nu-parser/src/lib.rs

@@ -12,6 +12,7 @@ pub use flatten::{flatten_block, flatten_expression, flatten_pipeline, FlatShape
 pub use known_external::KnownExternal;
 pub use lex::{lex, Token, TokenContents};
 pub use lite_parse::{lite_parse, LiteBlock};
+pub use parse_keywords::*;
 
 pub use parser::{
     is_math_expression_like, parse, parse_block, parse_duration_bytes, parse_external_call,

crates/nu-parser/src/parse_keywords.rs

@@ -1359,7 +1359,11 @@ pub fn parse_use(
                     }
                 } else {
                     error = error.or(Some(ParseError::ModuleNotFound(import_pattern.head.span)));
-                    (ImportPattern::new(), Overlay::new())
+
+                    let mut import_pattern = ImportPattern::new();
+                    import_pattern.head.span = spans[1];
+
+                    (import_pattern, Overlay::new())
                 }
             } else {
                 return (garbage_pipeline(spans), Some(ParseError::NonUtf8(spans[1])));