From 8688cdee11850ae4f08d11cdbc522dd106b1b8b2 Mon Sep 17 00:00:00 2001 From: blindfs Date: Tue, 22 Apr 2025 10:37:49 +0800 Subject: [PATCH 1/5] refactor(completion): custom_completion to `PositionalArg` and `Flag`, vibe coded by zed --- crates/nu-cli/src/completions/completer.rs | 28 +++++++++++++++-- crates/nu-parser/src/flatten.rs | 4 --- crates/nu-parser/src/parse_keywords.rs | 1 + crates/nu-parser/src/parser.rs | 36 ++++++++++++++++------ crates/nu-protocol/src/ast/expression.rs | 8 +---- crates/nu-protocol/src/signature.rs | 11 ++++++- crates/nu-protocol/tests/test_signature.rs | 5 +++ 7 files changed, 70 insertions(+), 23 deletions(-) diff --git a/crates/nu-cli/src/completions/completer.rs b/crates/nu-cli/src/completions/completer.rs index 2da8c35f20..00324eb74e 100644 --- a/crates/nu-cli/src/completions/completer.rs +++ b/crates/nu-cli/src/completions/completer.rs @@ -348,8 +348,32 @@ impl NuCompleter { for (arg_idx, arg) in call.arguments.iter().enumerate() { let span = arg.span(); if span.contains(pos) { - // if customized completion specified, it has highest priority - if let Some(decl_id) = arg.expr().and_then(|e| e.custom_completion) { + // Get custom completion from PositionalArg or Flag + let custom_completion_decl_id = { + // Check PositionalArg or Flag from Signature + let signature = working_set.get_decl(call.decl_id).signature(); + + match arg { + // For named arguments, check Flag + Argument::Named((name, short, _)) => { + // Try to find matching flag (long or short) + let flag = signature.get_long_flag(&name.item) + .or_else(|| short.as_ref().and_then(|s| + signature.get_short_flag(s.item.chars().next().unwrap_or('_')))); + flag.and_then(|f| f.custom_completion) + } + // For positional arguments, check PositionalArg + Argument::Positional(_) => { + // Find the right positional argument by index + let arg_pos = positional_arg_indices.len(); + signature.get_positional(arg_pos) + .and_then(|pos_arg| pos_arg.custom_completion) + } + _ => None + } + }; + + if let Some(decl_id) = custom_completion_decl_id { // for `--foo ` and `--foo=`, the arg span should be trimmed let (new_span, prefix) = if matches!(arg, Argument::Named(_)) { strip_placeholder_with_rsplit( diff --git a/crates/nu-parser/src/flatten.rs b/crates/nu-parser/src/flatten.rs index c49b1720df..54c590a511 100644 --- a/crates/nu-parser/src/flatten.rs +++ b/crates/nu-parser/src/flatten.rs @@ -183,10 +183,6 @@ fn flatten_expression_into( expr: &Expression, output: &mut Vec<(Span, FlatShape)>, ) { - if let Some(custom_completion) = &expr.custom_completion { - output.push((expr.span, FlatShape::Custom(*custom_completion))); - return; - } match &expr.expr { Expr::AttributeBlock(ab) => { diff --git a/crates/nu-parser/src/parse_keywords.rs b/crates/nu-parser/src/parse_keywords.rs index afa875e8be..ca1c617adc 100644 --- a/crates/nu-parser/src/parse_keywords.rs +++ b/crates/nu-parser/src/parse_keywords.rs @@ -350,6 +350,7 @@ pub fn parse_for(working_set: &mut StateWorkingSet, lite_command: &LiteCommand) shape: var_type.to_shape(), var_id: Some(*var_id), default_value: None, + custom_completion: None, }, ); } diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs index 06a38afe7e..6ea10ce2ab 100644 --- a/crates/nu-parser/src/parser.rs +++ b/crates/nu-parser/src/parser.rs @@ -1067,6 +1067,7 @@ pub fn parse_internal_call( desc: "".to_string(), var_id: None, default_value: None, + custom_completion: None, }) } @@ -1311,7 +1312,6 @@ pub fn parse_call(working_set: &mut StateWorkingSet, spans: &[Span], head: Span) span: _, span_id: _, ty, - custom_completion, } = &alias.clone().wrapped_call { trace!("parsing: alias of external call"); @@ -1325,14 +1325,13 @@ pub fn parse_call(working_set: &mut StateWorkingSet, spans: &[Span], head: Span) final_args.push(arg); } - let mut expression = Expression::new( + let expression = Expression::new( working_set, Expr::ExternalCall(head, final_args.into()), Span::concat(spans), ty.clone(), ); - expression.custom_completion = *custom_completion; return expression; } else { trace!("parsing: alias of internal call"); @@ -3631,6 +3630,7 @@ pub fn parse_row_condition(working_set: &mut StateWorkingSet, spans: &[Span]) -> shape: SyntaxShape::Any, var_id: Some(var_id), default_value: None, + custom_completion: None, }); compile_block(working_set, &mut block); @@ -3815,6 +3815,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> required: false, var_id: Some(var_id), default_value: None, + custom_completion: None, }, type_annotated: false, }); @@ -3876,6 +3877,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> required: false, var_id: Some(var_id), default_value: None, + custom_completion: None, }, type_annotated: false, }); @@ -3918,6 +3920,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> required: false, var_id: Some(var_id), default_value: None, + custom_completion: None, }, type_annotated: false, }); @@ -3985,6 +3988,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> shape: SyntaxShape::Any, var_id: Some(var_id), default_value: None, + custom_completion: None, }, required: false, type_annotated: false, @@ -4012,6 +4016,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> shape: SyntaxShape::Any, var_id: Some(var_id), default_value: None, + custom_completion: None, })); parse_mode = ParseMode::Arg; } @@ -4038,6 +4043,7 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> shape: SyntaxShape::Any, var_id: Some(var_id), default_value: None, + custom_completion: None, }, required: true, type_annotated: false, @@ -4056,22 +4062,30 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> //TODO check if we're replacing a custom parameter already match last { Arg::Positional { - arg: PositionalArg { shape, var_id, .. }, + arg: PositionalArg { shape, var_id, custom_completion, .. }, required: _, type_annotated, } => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type()); + // Extract custom_completion from CompleterWrapper if present + if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + *custom_completion = Some(*decl_id); + } *shape = syntax_shape; *type_annotated = true; } Arg::RestPositional(PositionalArg { - shape, var_id, .. + shape, var_id, custom_completion, .. }) => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), Type::List(Box::new(syntax_shape.to_type()))); + // Extract custom_completion from CompleterWrapper if present + if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + *custom_completion = Some(*decl_id); + } *shape = syntax_shape; } Arg::Flag { - flag: Flag { arg, var_id, .. }, + flag: Flag { arg, var_id, custom_completion, .. }, type_annotated, } => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type()); @@ -4082,6 +4096,10 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> span, )); } + // Extract custom_completion from CompleterWrapper if present + if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + *custom_completion = Some(*decl_id); + } *arg = Some(syntax_shape); *type_annotated = true; } @@ -5038,9 +5056,9 @@ pub fn parse_value( } match shape { - SyntaxShape::CompleterWrapper(shape, custom_completion) => { - let mut expression = parse_value(working_set, span, shape); - expression.custom_completion = Some(*custom_completion); + SyntaxShape::CompleterWrapper(shape, _custom_completion) => { + // Ignore the custom_completion field since it's now stored in PositionalArg/Flag + let expression = parse_value(working_set, span, shape); expression } SyntaxShape::Number => parse_number(working_set, span), diff --git a/crates/nu-protocol/src/ast/expression.rs b/crates/nu-protocol/src/ast/expression.rs index 5ffa3aab8b..ba7bd6d6e8 100644 --- a/crates/nu-protocol/src/ast/expression.rs +++ b/crates/nu-protocol/src/ast/expression.rs @@ -1,7 +1,7 @@ use crate::{ ast::{Argument, Block, Expr, ExternalArgument, ImportPattern, MatchPattern, RecordItem}, engine::StateWorkingSet, - BlockId, DeclId, GetSpan, Signature, Span, SpanId, Type, VarId, IN_VARIABLE_ID, + BlockId, GetSpan, Signature, Span, SpanId, Type, VarId, IN_VARIABLE_ID, }; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -15,7 +15,6 @@ pub struct Expression { pub span: Span, pub span_id: SpanId, pub ty: Type, - pub custom_completion: Option, } impl Expression { @@ -26,7 +25,6 @@ impl Expression { span, span_id, ty: Type::Any, - custom_completion: None, } } @@ -572,7 +570,6 @@ impl Expression { span, span_id, ty, - custom_completion: None, } } @@ -582,7 +579,6 @@ impl Expression { span, span_id, ty, - custom_completion: None, } } @@ -592,7 +588,6 @@ impl Expression { span, span_id: SpanId::new(0), ty, - custom_completion: None, } } @@ -602,7 +597,6 @@ impl Expression { span: self.span, span_id, ty: self.ty, - custom_completion: self.custom_completion, } } diff --git a/crates/nu-protocol/src/signature.rs b/crates/nu-protocol/src/signature.rs index 8405846eb2..ab59296f9f 100644 --- a/crates/nu-protocol/src/signature.rs +++ b/crates/nu-protocol/src/signature.rs @@ -1,6 +1,6 @@ use crate::{ engine::{Call, Command, CommandType, EngineState, Stack}, - BlockId, Example, PipelineData, ShellError, SyntaxShape, Type, Value, VarId, + BlockId, DeclId, Example, PipelineData, ShellError, SyntaxShape, Type, Value, VarId, }; use nu_derive_value::FromValue; use serde::{Deserialize, Serialize}; @@ -23,6 +23,7 @@ pub struct Flag { // For custom commands pub var_id: Option, pub default_value: Option, + pub custom_completion: Option, } /// The signature definition for a positional argument @@ -35,6 +36,7 @@ pub struct PositionalArg { // For custom commands pub var_id: Option, pub default_value: Option, + pub custom_completion: Option, } /// Command categories @@ -254,6 +256,7 @@ impl Signature { required: false, var_id: None, default_value: None, + custom_completion: None, }; self.named.push(flag); self @@ -318,6 +321,7 @@ impl Signature { shape: shape.into(), var_id: None, default_value: None, + custom_completion: None, }); self @@ -336,6 +340,7 @@ impl Signature { shape: shape.into(), var_id: None, default_value: None, + custom_completion: None, }); self @@ -353,6 +358,7 @@ impl Signature { shape: shape.into(), var_id: None, default_value: None, + custom_completion: None, }); self @@ -392,6 +398,7 @@ impl Signature { desc: desc.into(), var_id: None, default_value: None, + custom_completion: None, }); self @@ -415,6 +422,7 @@ impl Signature { desc: desc.into(), var_id: None, default_value: None, + custom_completion: None, }); self @@ -437,6 +445,7 @@ impl Signature { desc: desc.into(), var_id: None, default_value: None, + custom_completion: None, }); self diff --git a/crates/nu-protocol/tests/test_signature.rs b/crates/nu-protocol/tests/test_signature.rs index 011607959a..704c885347 100644 --- a/crates/nu-protocol/tests/test_signature.rs +++ b/crates/nu-protocol/tests/test_signature.rs @@ -45,6 +45,7 @@ fn test_signature_chained() { shape: SyntaxShape::String, var_id: None, default_value: None, + custom_completion: None, }) ); assert_eq!( @@ -55,6 +56,7 @@ fn test_signature_chained() { shape: SyntaxShape::String, var_id: None, default_value: None, + custom_completion: None, }) ); assert_eq!( @@ -65,6 +67,7 @@ fn test_signature_chained() { shape: SyntaxShape::String, var_id: None, default_value: None, + custom_completion: None, }) ); @@ -78,6 +81,7 @@ fn test_signature_chained() { desc: "required named description".to_string(), var_id: None, default_value: None, + custom_completion: None, }) ); @@ -91,6 +95,7 @@ fn test_signature_chained() { desc: "required named description".to_string(), var_id: None, default_value: None, + custom_completion: None, }) ); } From 50cb7623ac005e8111e871aadef960a429fe49a6 Mon Sep 17 00:00:00 2001 From: blindfs Date: Tue, 22 Apr 2025 10:57:36 +0800 Subject: [PATCH 2/5] fix: custom completion for flags only if we are on the flag value --- crates/nu-cli/src/completions/completer.rs | 31 +++++++++++++++------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/crates/nu-cli/src/completions/completer.rs b/crates/nu-cli/src/completions/completer.rs index 00324eb74e..dc3919e220 100644 --- a/crates/nu-cli/src/completions/completer.rs +++ b/crates/nu-cli/src/completions/completer.rs @@ -352,27 +352,38 @@ impl NuCompleter { let custom_completion_decl_id = { // Check PositionalArg or Flag from Signature let signature = working_set.get_decl(call.decl_id).signature(); - + match arg { // For named arguments, check Flag - Argument::Named((name, short, _)) => { - // Try to find matching flag (long or short) - let flag = signature.get_long_flag(&name.item) - .or_else(|| short.as_ref().and_then(|s| - signature.get_short_flag(s.item.chars().next().unwrap_or('_')))); - flag.and_then(|f| f.custom_completion) + Argument::Named((name, short, value)) => { + if value.as_ref().is_none_or(|e| !e.span.contains(pos)) { + None + } else { + // If we're completing the value of the flag, + // search for the matching custom completion decl_id (long or short) + let flag = + signature.get_long_flag(&name.item).or_else(|| { + short.as_ref().and_then(|s| { + signature.get_short_flag( + s.item.chars().next().unwrap_or('_'), + ) + }) + }); + flag.and_then(|f| f.custom_completion) + } } // For positional arguments, check PositionalArg Argument::Positional(_) => { // Find the right positional argument by index let arg_pos = positional_arg_indices.len(); - signature.get_positional(arg_pos) + signature + .get_positional(arg_pos) .and_then(|pos_arg| pos_arg.custom_completion) } - _ => None + _ => None, } }; - + if let Some(decl_id) = custom_completion_decl_id { // for `--foo ` and `--foo=`, the arg span should be trimmed let (new_span, prefix) = if matches!(arg, Argument::Named(_)) { From 27c3fd670cedfb1040e51c49ed9067e592fa2cec Mon Sep 17 00:00:00 2001 From: blindfs Date: Tue, 22 Apr 2025 11:16:05 +0800 Subject: [PATCH 3/5] test: vibe coded --- crates/nu-cli/tests/completions/mod.rs | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/crates/nu-cli/tests/completions/mod.rs b/crates/nu-cli/tests/completions/mod.rs index c8a02790dd..6402116e2b 100644 --- a/crates/nu-cli/tests/completions/mod.rs +++ b/crates/nu-cli/tests/completions/mod.rs @@ -2258,6 +2258,30 @@ fn extern_custom_completion_short_flag(mut extern_completer: NuCompleter) { match_suggestions(&expected, &suggestions); } +/// When we're completing the flag name itself, not its value, +/// custom completions should not be used +#[rstest] +fn custom_completion_flag_name_not_value(mut extern_completer: NuCompleter) { + let suggestions = extern_completer.complete("spam --f", 8); + assert!( + suggestions.iter().any(|s| s.value == "--foo"), + "Should contain --foo flag" + ); + let should_not_contain: Vec<_> = vec!["cat", "dog", "eel"]; + for item in should_not_contain { + assert!( + !suggestions.iter().any(|s| s.value == item), + "Should not contain custom completion {}", + item + ); + } + + // Also test with partial short flag + let suggestions = extern_completer.complete("spam -f", 7); + assert_eq!(1, suggestions.len(), "Should only have one suggestion"); + assert_eq!("-f", suggestions[0].value, "Should suggest -f flag"); +} + #[rstest] fn extern_complete_flags(mut extern_completer: NuCompleter) { let suggestions = extern_completer.complete("spam -", 6); From 5655bb62abd0523343f2b906573c9dc238f67685 Mon Sep 17 00:00:00 2001 From: blindfs Date: Tue, 22 Apr 2025 11:23:43 +0800 Subject: [PATCH 4/5] fix: clippy --- crates/nu-parser/src/parser.rs | 36 ++++++++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs index 6ea10ce2ab..07ff52729c 100644 --- a/crates/nu-parser/src/parser.rs +++ b/crates/nu-parser/src/parser.rs @@ -4062,30 +4062,49 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> //TODO check if we're replacing a custom parameter already match last { Arg::Positional { - arg: PositionalArg { shape, var_id, custom_completion, .. }, + arg: + PositionalArg { + shape, + var_id, + custom_completion, + .. + }, required: _, type_annotated, } => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type()); // Extract custom_completion from CompleterWrapper if present - if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + if let SyntaxShape::CompleterWrapper(_, decl_id) = + &syntax_shape + { *custom_completion = Some(*decl_id); } *shape = syntax_shape; *type_annotated = true; } Arg::RestPositional(PositionalArg { - shape, var_id, custom_completion, .. + shape, + var_id, + custom_completion, + .. }) => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), Type::List(Box::new(syntax_shape.to_type()))); // Extract custom_completion from CompleterWrapper if present - if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + if let SyntaxShape::CompleterWrapper(_, decl_id) = + &syntax_shape + { *custom_completion = Some(*decl_id); } *shape = syntax_shape; } Arg::Flag { - flag: Flag { arg, var_id, custom_completion, .. }, + flag: + Flag { + arg, + var_id, + custom_completion, + .. + }, type_annotated, } => { working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type()); @@ -4097,7 +4116,9 @@ pub fn parse_signature_helper(working_set: &mut StateWorkingSet, span: Span) -> )); } // Extract custom_completion from CompleterWrapper if present - if let SyntaxShape::CompleterWrapper(_, decl_id) = &syntax_shape { + if let SyntaxShape::CompleterWrapper(_, decl_id) = + &syntax_shape + { *custom_completion = Some(*decl_id); } *arg = Some(syntax_shape); @@ -5058,8 +5079,7 @@ pub fn parse_value( match shape { SyntaxShape::CompleterWrapper(shape, _custom_completion) => { // Ignore the custom_completion field since it's now stored in PositionalArg/Flag - let expression = parse_value(working_set, span, shape); - expression + parse_value(working_set, span, shape) } SyntaxShape::Number => parse_number(working_set, span), SyntaxShape::Float => parse_float(working_set, span), From b17a46870cc85468c2b57a08519de0cc4ed15fb5 Mon Sep 17 00:00:00 2001 From: blindfs Date: Tue, 22 Apr 2025 11:27:08 +0800 Subject: [PATCH 5/5] fix: fmt --- crates/nu-parser/src/flatten.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/nu-parser/src/flatten.rs b/crates/nu-parser/src/flatten.rs index 54c590a511..6f8016062e 100644 --- a/crates/nu-parser/src/flatten.rs +++ b/crates/nu-parser/src/flatten.rs @@ -183,7 +183,6 @@ fn flatten_expression_into( expr: &Expression, output: &mut Vec<(Span, FlatShape)>, ) { - match &expr.expr { Expr::AttributeBlock(ab) => { for attr in &ab.attributes {