From 61f598d8a733eba35810a4de6dc76f17433bbefc Mon Sep 17 00:00:00 2001
From: Dongbo Wang <dongbow@microsoft.com>
Date: Tue, 14 May 2024 13:34:37 -0700
Subject: [PATCH] Update `SelectCommandArgument` to properly handle POSIX style
 options for CLI commands (#4016)

---
 PSReadLine/KillYank.cs | 24 +++++++++++++++++++-
 test/KillYankTest.cs   | 50 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/PSReadLine/KillYank.cs b/PSReadLine/KillYank.cs
index f05fc8b..77fe33f 100644
--- a/PSReadLine/KillYank.cs
+++ b/PSReadLine/KillYank.cs
@@ -6,6 +6,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Management.Automation.Language;
+using System.Text.RegularExpressions;
 using Microsoft.PowerShell.Internal;
 
 namespace Microsoft.PowerShell
@@ -29,6 +30,10 @@ namespace Microsoft.PowerShell
         private YankLastArgState _yankLastArgState;
         private int _visualSelectionCommandCount;
 
+        // Pattern to check for CLI parameters like '--json'.
+        // Valid characters are 'a-z', 'A-Z', '0-9', '_' (all covered by '\w'), and '-'.
+        private static readonly Regex s_cliOptionPattern = new(@"^--[\w-]+$", RegexOptions.Compiled);
+
         /// <summary>
         /// Mark the current location of the cursor for use in a subsequent editing command.
         /// </summary>
@@ -480,7 +485,7 @@ namespace Microsoft.PowerShell
                         var argument = cmdAst.CommandElements[j] switch
                         {
                             CommandParameterAst paramAst => paramAst.Argument,
-                            ExpressionAst expAst => expAst,
+                            ExpressionAst exprAst => ProcessExpressionAst(exprAst),
                             _ => null,
                         };
 
@@ -614,6 +619,7 @@ namespace Microsoft.PowerShell
             _singleton.VisualSelectionCommon(() => SetCursorPosition(newEndCursor), forceSetMark: true);
 
 
+            // =====  Local Functions  =====
             // Get the script block AST's whose extent contains the cursor.
             bool GetScriptBlockAst(Ast ast)
             {
@@ -639,6 +645,22 @@ namespace Microsoft.PowerShell
                     ? ast.Extent.EndOffset - 1 > cursor
                     : ast.Extent.EndOffset >= cursor;
             }
+
+            // Process an expression AST to check if it's a CLI posix style option.
+            static ExpressionAst ProcessExpressionAst(ExpressionAst exprAst)
+            {
+                if (exprAst is StringConstantExpressionAst strAst
+                    && strAst.StringConstantType is StringConstantType.BareWord
+                    && strAst.Value.StartsWith("--")
+                    && s_cliOptionPattern.IsMatch(strAst.Value))
+                {
+                    // It's a CLI posix style option, like '--json' or '--machine-type',
+                    // so we treat it as a parameter.
+                    return null;
+                }
+
+                return exprAst;
+            }
         }
 
         /// <summary>
diff --git a/test/KillYankTest.cs b/test/KillYankTest.cs
index 3433f7d..8705faa 100644
--- a/test/KillYankTest.cs
+++ b/test/KillYankTest.cs
@@ -619,6 +619,56 @@ namespace Test
                 _.Escape));
         }
 
+        [SkippableFact]
+        public void SelectCommandArgument_CLIArgs()
+        {
+            TestSetup(KeyMode.Cmd);
+
+            Test("", Keys(
+                "az webapp --name MyWebApp --resource-group MyResourceGroup",
+                _.Alt_a, CheckThat(() => AssertScreenIs(1,
+                    TokenClassification.Command, "az",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Selection, "webapp",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Parameter, "--name",
+                    TokenClassification.None, " MyWebApp ",
+                    TokenClassification.Parameter, "--resource-group",
+                    TokenClassification.None, " MyResourceGroup ")),
+
+                _.Alt_a, CheckThat(() => AssertScreenIs(1,
+                    TokenClassification.Command, "az",
+                    TokenClassification.None, " webapp ",
+                    TokenClassification.Parameter, "--name",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Selection, "MyWebApp",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Parameter, "--resource-group",
+                    TokenClassification.None, " MyResourceGroup ")),
+
+                _.Alt_a, CheckThat(() => AssertScreenIs(1,
+                    TokenClassification.Command, "az",
+                    TokenClassification.None, " webapp ",
+                    TokenClassification.Parameter, "--name",
+                    TokenClassification.None, " MyWebApp ",
+                    TokenClassification.Parameter, "--resource-group",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Selection, "MyResourceGroup")),
+
+                // Verify that we can loop through the arguments.
+                _.Alt_a, CheckThat(() => AssertScreenIs(1,
+                    TokenClassification.Command, "az",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Selection, "webapp",
+                    TokenClassification.None, ' ',
+                    TokenClassification.Parameter, "--name",
+                    TokenClassification.None, " MyWebApp ",
+                    TokenClassification.Parameter, "--resource-group",
+                    TokenClassification.None, " MyResourceGroup ")),
+
+                _.Escape));
+        }
+
         [SkippableFact]
         public void SelectCommandArgument_HereStringArgs()
         {