Improve data.table column names autocomplete. (#716)

Author: dfalbel

crates/ark/src/lsp/completions/sources/composite/subset.rs

@@ -68,14 +68,15 @@ pub(super) fn completions_from_subset(
 
     let text = context.document.contents.node_slice(&child)?.to_string();
 
-    completions_from_evaluated_object_names(&text, ENQUOTE)
+    completions_from_evaluated_object_names(&text, ENQUOTE, node.node_type())
 }
 
 #[cfg(test)]
 mod tests {
     use harp::eval::RParseEvalOptions;
     use tree_sitter::Point;
 
+    use crate::fixtures::package_is_installed;
     use crate::lsp::completions::sources::composite::subset::completions_from_subset;
     use crate::lsp::document_context::DocumentContext;
     use crate::lsp::documents::Document;
@@ -126,4 +127,66 @@ mod tests {
             harp::parse_eval("remove(foo)", options.clone()).unwrap();
         })
     }
+
+    #[test]
+    fn test_data_table_subset_completions() {
+        r_task(|| {
+            if !package_is_installed("data.table") {
+                return;
+            }
+
+            harp::parse_eval_global("x <- data.table::as.data.table(mtcars)").unwrap();
+
+            // Works for single comlumn name completion
+            let point = Point { row: 0, column: 2 };
+            let document = Document::new("x[]", None);
+            let context = DocumentContext::new(&document, point, None);
+
+            let completions = completions_from_subset(&context).unwrap().unwrap();
+            assert_eq!(completions.len(), 11);
+
+            let completion = completions.get(0).unwrap();
+            assert_eq!(completion.label, "mpg".to_string());
+            assert_eq!(completion.insert_text, None); // No enquote, means the label is used directly
+
+            // Works when completing inside a `c` call
+            let point = Point { row: 0, column: 6 };
+            let document = Document::new("x[, c()]", None);
+            let context = DocumentContext::new(&document, point, None);
+
+            let completions = completions_from_subset(&context).unwrap().unwrap();
+            assert_eq!(completions.len(), 11);
+
+            let completion = completions.get(0).unwrap();
+            assert_eq!(completion.label, "mpg".to_string());
+            assert_eq!(completion.insert_text, None); // No enquote, means the label is used directly
+
+            // Works when completing inside a `c` call with a collumn already selected
+            let point = Point { row: 0, column: 10 };
+            let document = Document::new("x[, c(mpg,)]", None);
+            let context = DocumentContext::new(&document, point, None);
+
+            let completions = completions_from_subset(&context).unwrap().unwrap();
+            assert_eq!(completions.len(), 11);
+
+            let completion = completions.get(0).unwrap();
+            // TODO: ideally we could assert that mpg doesn't appear again, or appears at the end
+            assert_eq!(completion.label, "mpg".to_string());
+            assert_eq!(completion.insert_text, None); // No enquote, means the label is used directly
+
+            // Works completing subset2
+            let point = Point { row: 0, column: 3 };
+            let document = Document::new("x[[]]", None);
+            let context = DocumentContext::new(&document, point, None);
+
+            let completions = completions_from_subset(&context).unwrap().unwrap();
+            assert_eq!(completions.len(), 11);
+
+            let completion = completions.get(0).unwrap();
+            assert_eq!(completion.label, "mpg".to_string());
+            assert_eq!(completion.insert_text, Some("\"mpg\"".to_string())); // Enquoted result
+
+            harp::parse_eval_global("remove(x)").unwrap();
+        })
+    }
 }

crates/ark/src/lsp/completions/sources/unique/subset.rs

@@ -50,7 +50,9 @@ pub(super) fn completions_from_string_subset(
 
     let text = context.document.contents.node_slice(&node)?.to_string();
 
-    if let Some(mut candidates) = completions_from_evaluated_object_names(&text, ENQUOTE)? {
+    if let Some(mut candidates) =
+        completions_from_evaluated_object_names(&text, ENQUOTE, node.node_type())?
+    {
         completions.append(&mut candidates);
     }
 

crates/ark/src/lsp/completions/sources/utils.rs

@@ -11,6 +11,7 @@ use harp::eval::RParseEvalOptions;
 use harp::exec::RFunction;
 use harp::exec::RFunctionExt;
 use harp::object::RObject;
+use harp::utils::r_inherits;
 use regex::Regex;
 use tower_lsp::lsp_types::CompletionItem;
 use tree_sitter::Node;
@@ -170,6 +171,7 @@ fn call_prev_leaf_position_type(node: &Node, allow_ambiguous: bool) -> CallNodeP
 pub(super) fn completions_from_evaluated_object_names(
     name: &str,
     enquote: bool,
+    node_type: NodeType,
 ) -> Result<Option<Vec<CompletionItem>>> {
     log::info!("completions_from_evaluated_object_names({name:?})");
 
@@ -204,6 +206,15 @@ pub(super) fn completions_from_evaluated_object_names(
     let completions = if harp::utils::r_is_matrix(object.sexp) {
         // Special case just for 2D arrays
         completions_from_object_colnames(object, name, enquote)?
+    } else if r_inherits(object.sexp, "data.table") {
+        // The `[` method for data.table uses NSE so we don't enquote the names
+        // https://github.com/posit-dev/positron/issues/3140
+        let enquote = match node_type {
+            NodeType::Subset => false,
+            NodeType::Subset2 => true,
+            _ => enquote,
+        };
+        completions_from_object_names(object, name, enquote)?
     } else {
         completions_from_object_names(object, name, enquote)?
     };
@@ -259,6 +270,7 @@ mod tests {
     use harp::eval::parse_eval_global;
     use tree_sitter::Point;
 
+    use crate::fixtures::package_is_installed;
     use crate::lsp::completions::sources::utils::call_node_position_type;
     use crate::lsp::completions::sources::utils::completions_from_evaluated_object_names;
     use crate::lsp::completions::sources::utils::CallNodePositionType;
@@ -426,7 +438,7 @@ mod tests {
             parse_eval_global("x <- 1:2").unwrap();
             parse_eval_global("names(x) <- c('a', 'b')").unwrap();
 
-            let completions = completions_from_evaluated_object_names("x", false)
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
                 .unwrap()
                 .unwrap();
             assert_eq!(completions.len(), 2);
@@ -438,7 +450,7 @@ mod tests {
             // Data frame
             parse_eval_global("x <- data.frame(a = 1, b = 2, c = 3)").unwrap();
 
-            let completions = completions_from_evaluated_object_names("x", false)
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
                 .unwrap()
                 .unwrap();
             assert_eq!(completions.len(), 3);
@@ -452,7 +464,7 @@ mod tests {
             parse_eval_global("x <- array(1:2)").unwrap();
             parse_eval_global("names(x) <- c('a', 'b')").unwrap();
 
-            let completions = completions_from_evaluated_object_names("x", false)
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
                 .unwrap()
                 .unwrap();
             assert_eq!(completions.len(), 2);
@@ -466,7 +478,7 @@ mod tests {
             parse_eval_global("rownames(x) <- 'a'").unwrap();
             parse_eval_global("colnames(x) <- 'b'").unwrap();
 
-            let completions = completions_from_evaluated_object_names("x", false)
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
                 .unwrap()
                 .unwrap();
             assert_eq!(completions.len(), 1);
@@ -482,12 +494,48 @@ mod tests {
             parse_eval_global("rownames(x) <- 'a'").unwrap();
             parse_eval_global("colnames(x) <- 'b'").unwrap();
 
-            let completions = completions_from_evaluated_object_names("x", false)
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
                 .unwrap()
                 .unwrap();
             assert!(completions.is_empty());
 
             parse_eval_global("remove(x)").unwrap();
         })
     }
+
+    #[test]
+    fn test_data_table_completions() {
+        r_task(|| {
+            // Skip test if data.table is not installed
+            if !package_is_installed("data.table") {
+                return;
+            }
+
+            parse_eval_global("x <- data.table::as.data.table(mtcars)").unwrap();
+
+            // Subset completions
+            let completions = completions_from_evaluated_object_names("x", false, NodeType::Subset)
+                .unwrap()
+                .unwrap();
+
+            assert_eq!(completions.len(), 11);
+            assert_eq!(completions.get(0).unwrap().label, String::from("mpg"));
+            assert_eq!(completions.get(0).unwrap().insert_text, None);
+
+            // Subset2 completions
+            let completions =
+                completions_from_evaluated_object_names("x", false, NodeType::Subset2)
+                    .unwrap()
+                    .unwrap();
+
+            assert_eq!(completions.len(), 11);
+            assert_eq!(completions.get(0).unwrap().label, String::from("mpg"));
+            assert_eq!(
+                completions.get(0).unwrap().insert_text,
+                Some("\"mpg\"".to_string())
+            );
+
+            parse_eval_global("remove(x)").unwrap();
+        })
+    }
 }