Add completion for variables assigned in ArrayLiterals and ParenExpressions (#25303)

MartinGC94 · web-flow · commit 957eb1b2791f · 2025-05-08T18:17:40.000+05:00
diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
@@ -5573,7 +5573,13 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a
                         : AstVisitAction.StopVisit;
                 }
 
-                if (assignmentStatementAst.Left is AttributedExpressionAst attributedExpression)
+                ProcessAssignmentLeftSide(assignmentStatementAst.Left, assignmentStatementAst.Right);
+                return AstVisitAction.Continue;
+            }
+
+            private void ProcessAssignmentLeftSide(ExpressionAst left, StatementAst right)
+            {
+                if (left is AttributedExpressionAst attributedExpression)
                 {
                     var firstConvertExpression = attributedExpression as ConvertExpressionAst;
                     ExpressionAst child = attributedExpression.Child;
@@ -5593,7 +5599,7 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a
                     {
                         if (variableExpression == CompletionVariableAst || s_specialVariablesCache.Value.Contains(variableExpression.VariablePath.UserPath))
                         {
-                            return AstVisitAction.Continue;
+                            return;
                         }
 
                         if (firstConvertExpression is not null)
@@ -5602,22 +5608,22 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a
                         }
                         else
                         {
-                            PSTypeName lastAssignedType = assignmentStatementAst.Right is CommandExpressionAst commandExpression
+                            PSTypeName lastAssignedType = right is CommandExpressionAst commandExpression
                                 ? GetInferredVarTypeFromAst(commandExpression.Expression)
                                 : null;
                             SaveVariableInfo(variableExpression.VariablePath.UnqualifiedPath, lastAssignedType, isConstraint: false);
                         }
                     }
                 }
-                else if (assignmentStatementAst.Left is VariableExpressionAst variableExpression)
+                else if (left is VariableExpressionAst variableExpression)
                 {
                     if (variableExpression == CompletionVariableAst || s_specialVariablesCache.Value.Contains(variableExpression.VariablePath.UserPath))
                     {
-                        return AstVisitAction.Continue;
+                        return;
                     }
 
                     PSTypeName lastAssignedType;
-                    if (assignmentStatementAst.Right is CommandExpressionAst commandExpression)
+                    if (right is CommandExpressionAst commandExpression)
                     {
                         lastAssignedType = GetInferredVarTypeFromAst(commandExpression.Expression);
                     }
@@ -5628,8 +5634,21 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a
 
                     SaveVariableInfo(variableExpression.VariablePath.UnqualifiedPath, lastAssignedType, isConstraint: false);
                 }
-
-                return AstVisitAction.Continue;
+                else if (left is ArrayLiteralAst array)
+                {
+                    foreach (ExpressionAst expression in array.Elements)
+                    {
+                        ProcessAssignmentLeftSide(expression, right);
+                    }
+                }
+                else if (left is ParenExpressionAst parenExpression)
+                {
+                    ExpressionAst pureExpression = parenExpression.Pipeline.GetPureExpression();
+                    if (pureExpression is not null)
+                    {
+                        ProcessAssignmentLeftSide(pureExpression, right);
+                    }
+                }
             }
 
             public override AstVisitAction VisitCommand(CommandAst commandAst)
diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
@@ -1084,6 +1084,17 @@ param([ValidatePattern(
         $res.CompletionMatches[0].CompletionText | Should -BeExactly '$TestVar1'
     }
 
+    It 'Should complete variable assigned in ParenExpression' {
+        $res = TabExpansion2 -inputScript '($ParenVar) = 1; $ParenVa'
+        $res.CompletionMatches[0].CompletionText | Should -BeExactly '$ParenVar'
+    }
+
+    It 'Should complete variable assigned in ArrayLiteral' {
+        $res = TabExpansion2 -inputScript '$DemoVar1, $DemoVar2 = 1..10; $DemoVar'
+        $res.CompletionMatches[0].CompletionText | Should -BeExactly '$DemoVar1'
+        $res.CompletionMatches[1].CompletionText | Should -BeExactly '$DemoVar2'
+    }
+
     Context 'Start-Process -Verb parameter completion' {
         BeforeAll {
             function GetProcessInfoVerbs([string]$path, [switch]$singleQuote, [switch]$doubleQuote) {

Original file line number	Diff line number	Diff line change
`@@ -5573,7 +5573,13 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a`
`5573`	`5573`	`: AstVisitAction.StopVisit;`
`5574`	`5574`	`}`
`5575`	`5575`
`5576`		`- if (assignmentStatementAst.Left is AttributedExpressionAst attributedExpression)`
	`5576`	`+ ProcessAssignmentLeftSide(assignmentStatementAst.Left, assignmentStatementAst.Right);`
	`5577`	`+ return AstVisitAction.Continue;`
	`5578`	`+ }`
	`5579`	`+`
	`5580`	`+ private void ProcessAssignmentLeftSide(ExpressionAst left, StatementAst right)`
	`5581`	`+ {`
	`5582`	`+ if (left is AttributedExpressionAst attributedExpression)`
`5577`	`5583`	`{`
`5578`	`5584`	`var firstConvertExpression = attributedExpression as ConvertExpressionAst;`
`5579`	`5585`	`ExpressionAst child = attributedExpression.Child;`
`@@ -5593,7 +5599,7 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a`
`5593`	`5599`	`{`
`5594`	`5600`	`if (variableExpression == CompletionVariableAst \|\| s_specialVariablesCache.Value.Contains(variableExpression.VariablePath.UserPath))`
`5595`	`5601`	`{`
`5596`		`- return AstVisitAction.Continue;`
	`5602`	`+ return;`
`5597`	`5603`	`}`
`5598`	`5604`
`5599`	`5605`	`if (firstConvertExpression is not null)`
`@@ -5602,22 +5608,22 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a`
`5602`	`5608`	`}`
`5603`	`5609`	`else`
`5604`	`5610`	`{`
`5605`		`- PSTypeName lastAssignedType = assignmentStatementAst.Right is CommandExpressionAst commandExpression`
	`5611`	`+ PSTypeName lastAssignedType = right is CommandExpressionAst commandExpression`
`5606`	`5612`	`? GetInferredVarTypeFromAst(commandExpression.Expression)`
`5607`	`5613`	`: null;`
`5608`	`5614`	`SaveVariableInfo(variableExpression.VariablePath.UnqualifiedPath, lastAssignedType, isConstraint: false);`
`5609`	`5615`	`}`
`5610`	`5616`	`}`
`5611`	`5617`	`}`
`5612`		`- else if (assignmentStatementAst.Left is VariableExpressionAst variableExpression)`
	`5618`	`+ else if (left is VariableExpressionAst variableExpression)`
`5613`	`5619`	`{`
`5614`	`5620`	`if (variableExpression == CompletionVariableAst \|\| s_specialVariablesCache.Value.Contains(variableExpression.VariablePath.UserPath))`
`5615`	`5621`	`{`
`5616`		`- return AstVisitAction.Continue;`
	`5622`	`+ return;`
`5617`	`5623`	`}`
`5618`	`5624`
`5619`	`5625`	`PSTypeName lastAssignedType;`
`5620`		`- if (assignmentStatementAst.Right is CommandExpressionAst commandExpression)`
	`5626`	`+ if (right is CommandExpressionAst commandExpression)`
`5621`	`5627`	`{`
`5622`	`5628`	`lastAssignedType = GetInferredVarTypeFromAst(commandExpression.Expression);`
`5623`	`5629`	`}`
`@@ -5628,8 +5634,21 @@ public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst a`
`5628`	`5634`
`5629`	`5635`	`SaveVariableInfo(variableExpression.VariablePath.UnqualifiedPath, lastAssignedType, isConstraint: false);`
`5630`	`5636`	`}`
`5631`		`-`
`5632`		`- return AstVisitAction.Continue;`
	`5637`	`+ else if (left is ArrayLiteralAst array)`
	`5638`	`+ {`
	`5639`	`+ foreach (ExpressionAst expression in array.Elements)`
	`5640`	`+ {`
	`5641`	`+ ProcessAssignmentLeftSide(expression, right);`
	`5642`	`+ }`
	`5643`	`+ }`
	`5644`	`+ else if (left is ParenExpressionAst parenExpression)`
	`5645`	`+ {`
	`5646`	`+ ExpressionAst pureExpression = parenExpression.Pipeline.GetPureExpression();`
	`5647`	`+ if (pureExpression is not null)`
	`5648`	`+ {`
	`5649`	`+ ProcessAssignmentLeftSide(pureExpression, right);`
	`5650`	`+ }`
	`5651`	`+ }`
`5633`	`5652`	`}`
`5634`	`5653`
`5635`	`5654`	`public override AstVisitAction VisitCommand(CommandAst commandAst)`