nushell · sholderbach · Apr 15, 2023 · Apr 8, 2023 · Apr 8, 2023 · Apr 14, 2023
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
@@ -455,19 +455,19 @@ fn parse_short_flags(
     if let Ok(arg_contents_uft8_ref) = str::from_utf8(arg_contents) {
         if arg_contents_uft8_ref.starts_with('-') && arg_contents_uft8_ref.len() > 1 {
             let short_flags = &arg_contents_uft8_ref[1..];
+            let num_chars = short_flags.chars().count();
             let mut found_short_flags = vec![];
             let mut unmatched_short_flags = vec![];
-            for short_flag in short_flags.char_indices() {
-                let short_flag_char = short_flag.1;
-                let orig = arg_span;
+            for (offset, short_flag) in short_flags.char_indices() {
                 let short_flag_span = Span::new(
-                    orig.start + 1 + short_flag.0,
-                    orig.start + 1 + short_flag.0 + short_flag_char.len_utf8(),
+                    arg_span.start + 1 + offset,
+                    arg_span.start + 1 + offset + short_flag.len_utf8(),
                 );
-                if let Some(flag) = sig.get_short_flag(short_flag_char) {
-                    // If we require an arg and are in a batch of short flags, error
-                    if !found_short_flags.is_empty() && flag.arg.is_some() {
-                        working_set.error(ParseError::ShortFlagBatchCantTakeArg(short_flag_span));
+                if let Some(flag) = sig.get_short_flag(short_flag) {
+                    // Allow args in short flag batches as long as it is the last flag.
+                    if flag.arg.is_some() && offset < num_chars - 1 {
+                        working_set
+                            .error(ParseError::OnlyLastFlagInBatchCanTakeArg(short_flag_span));
                         break;
                     }
                     found_short_flags.push(flag);

diff --git a/crates/nu-parser/tests/test_parser.rs b/crates/nu-parser/tests/test_parser.rs
@@ -1,5 +1,5 @@
 use nu_parser::*;
-use nu_protocol::ast::{Call, PathMember};
+use nu_protocol::ast::{Argument, Call, PathMember};
 use nu_protocol::Span;
 use nu_protocol::{
     ast::{Expr, Expression, PipelineElement},
@@ -638,18 +638,68 @@ pub fn parse_call_missing_short_flag_arg() {
 }
 
 #[test]
-pub fn parse_call_too_many_shortflag_args() {
+pub fn parse_call_short_flag_batch_arg_allowed() {
     let engine_state = EngineState::new();
     let mut working_set = StateWorkingSet::new(&engine_state);
 
     let sig = Signature::build("foo")
         .named("--jazz", SyntaxShape::Int, "jazz!!", Some('j'))
-        .named("--math", SyntaxShape::Int, "math!!", Some('m'));
+        .switch("--math", "math!!", Some('m'));
     working_set.add_decl(sig.predeclare());
-    parse(&mut working_set, None, b"foo -mj", true);
+
+    let block = parse(&mut working_set, None, b"foo -mj 10", true);
+
+    assert!(working_set.parse_errors.is_empty());
+    assert_eq!(block.len(), 1);
+    let expressions = &block[0];
+    assert_eq!(expressions.len(), 1);
+
+    if let PipelineElement::Expression(
+        _,
+        Expression {
+            expr: Expr::Call(call),
+            ..
+        },
+    ) = &expressions[0]
+    {
+        assert_eq!(call.decl_id, 0);
+        assert_eq!(call.arguments.len(), 2);
+        matches!(call.arguments[0], Argument::Named((_, None, None)));
+        matches!(call.arguments[1], Argument::Named((_, None, Some(_))));
+    }
+}
+
+#[test]
+pub fn parse_call_short_flag_batch_arg_disallowed() {
+    let engine_state = EngineState::new();
+    let mut working_set = StateWorkingSet::new(&engine_state);
+
+    let sig = Signature::build("foo")
+        .named("--jazz", SyntaxShape::Int, "jazz!!", Some('j'))
+        .switch("--math", "math!!", Some('m'));
+    working_set.add_decl(sig.predeclare());
+
+    parse(&mut working_set, None, b"foo -jm 10", true);
+    assert!(matches!(
+        working_set.parse_errors.first(),
+        Some(ParseError::OnlyLastFlagInBatchCanTakeArg(..))
+    ));
+}
+
+#[test]
+pub fn parse_call_short_flag_batch_disallow_multiple_args() {
+    let engine_state = EngineState::new();
+    let mut working_set = StateWorkingSet::new(&engine_state);
+
+    let sig = Signature::build("foo")
+        .named("--math", SyntaxShape::Int, "math!!", Some('m'))
+        .named("--jazz", SyntaxShape::Int, "jazz!!", Some('j'));
+    working_set.add_decl(sig.predeclare());
+
+    parse(&mut working_set, None, b"foo -mj 10 20", true);
     assert!(matches!(
         working_set.parse_errors.first(),
-        Some(ParseError::ShortFlagBatchCantTakeArg(..))
+        Some(ParseError::OnlyLastFlagInBatchCanTakeArg(..))
     ));
 }
 

diff --git a/crates/nu-protocol/src/parse_error.rs b/crates/nu-protocol/src/parse_error.rs
@@ -292,9 +292,9 @@ pub enum ParseError {
     #[diagnostic(code(nu::parser::missing_flag_param))]
     MissingFlagParam(String, #[label = "flag missing {0} argument"] Span),
 
-    #[error("Batches of short flags can't take arguments.")]
-    #[diagnostic(code(nu::parser::short_flag_arg_cant_take_arg))]
-    ShortFlagBatchCantTakeArg(#[label = "short flag batches can't take args"] Span),
+    #[error("Only the last flag in a short flag batch can take an argument.")]
+    #[diagnostic(code(nu::parser::only_last_flag_in_batch_can_take_arg))]
+    OnlyLastFlagInBatchCanTakeArg(#[label = "only the last flag can take args"] Span),
 
     #[error("Missing required positional argument.")]
     #[diagnostic(code(nu::parser::missing_positional), help("Usage: {2}"))]
@@ -473,7 +473,7 @@ impl ParseError {
             ParseError::RequiredAfterOptional(_, s) => *s,
             ParseError::UnknownType(s) => *s,
             ParseError::MissingFlagParam(_, s) => *s,
-            ParseError::ShortFlagBatchCantTakeArg(s) => *s,
+            ParseError::OnlyLastFlagInBatchCanTakeArg(s) => *s,
             ParseError::MissingPositional(_, s, _) => *s,
             ParseError::KeywordMissingArgument(_, _, s) => *s,
             ParseError::MissingType(s) => *s,

diff --git a/src/tests/test_known_external.rs b/src/tests/test_known_external.rs
@@ -40,10 +40,23 @@ fn known_external_complex_unknown_args() -> TestResult {
 }
 
 #[test]
-fn known_external_batched_short_flag_arg_disallowed() -> TestResult {
+fn known_external_short_flag_batch_arg_allowed() -> TestResult {
+    run_test_contains("extern echo [-a, -b: int]; echo -ab 10", "-b 10")
+}
+
+#[test]
+fn known_external_short_flag_batch_arg_disallowed() -> TestResult {
+    fail_test(
+        "extern echo [-a: int, -b]; echo -ab 10",
+        "last flag can take args",
+    )
+}
+
+#[test]
+fn known_external_short_flag_batch_multiple_args() -> TestResult {
     fail_test(
-        "extern echo [-a, -b: int]; echo -ab 10",
-        "short flag batches",
+        "extern echo [-a: int, -b: int]; echo -ab 10 20",
+        "last flag can take args",
     )
 }