diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionAnalysis.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionAnalysis.cs index b4e408ec22..5bfb5c9697 100644 --- a/src/System.Management.Automation/engine/CommandCompletion/CompletionAnalysis.cs +++ b/src/System.Management.Automation/engine/CommandCompletion/CompletionAnalysis.cs @@ -420,6 +420,17 @@ namespace System.Management.Automation case TokenKind.StringExpandable: case TokenKind.StringLiteral: + // Search to see if we're looking at an assignment + if (lastAst.Parent is CommandExpressionAst + && lastAst.Parent.Parent is AssignmentStatementAst assignmentAst) + { + // Handle scenarios like `$ErrorActionPreference = '` + if (TryGetCompletionsForVariableAssignment(completionContext, assignmentAst, out List completions)) + { + return completions; + } + } + result = GetResultForString(completionContext, ref replacementIndex, ref replacementLength, isQuotedString); break; @@ -505,6 +516,15 @@ namespace System.Management.Automation replacementIndex = completionContext.ReplacementIndex; replacementLength = completionContext.ReplacementLength; } + else if (lastAst.Parent is CommandExpressionAst + && lastAst.Parent.Parent is AssignmentStatementAst assignmentAst2) + { + // Handle scenarios like '[ValidateSet(11,22)][int]$i = 11; $i = 2' + if (TryGetCompletionsForVariableAssignment(completionContext, assignmentAst2, out List completions)) + { + result = completions; + } + } break; @@ -564,6 +584,17 @@ namespace System.Management.Automation // Handle scenarios such as 'gci | Format-Table @{Label=' if incomplete parsing of the assignment. return null; } + else if (lastAst is AssignmentStatementAst assignmentAst2) + { + completionContext.ReplacementIndex = replacementIndex += tokenAtCursor.Text.Length; + completionContext.ReplacementLength = replacementLength = 0; + + // Handle scenarios like '$ErrorActionPreference =' + if (TryGetCompletionsForVariableAssignment(completionContext, assignmentAst2, out List completions)) + { + return completions; + } + } else { // Handle scenarios such as 'configuration foo { File ab { Attributes =' @@ -732,6 +763,15 @@ namespace System.Management.Automation case TokenKind.Comma: case TokenKind.AtParen: { + if (lastAst is AssignmentStatementAst assignmentAst) + { + // Handle scenarios like '$ErrorActionPreference = ' + if (TryGetCompletionsForVariableAssignment(completionContext, assignmentAst, out result)) + { + break; + } + } + bool unused; result = GetResultForEnumPropertyValueOfDSCResource(completionContext, string.Empty, ref replacementIndex, ref replacementLength, out unused); break; @@ -1025,6 +1065,239 @@ namespace System.Management.Automation return keyValuePairWithCursor; } + // Pulls the variable out of an assignment's LHS expression + // Also brings back the innermost type constraint if there is one + private static VariableExpressionAst GetVariableFromExpressionAst( + ExpressionAst expression, + ref Type typeConstraint, + ref ValidateSetAttribute setConstraint) + { + switch (expression) + { + // $x = ... + case VariableExpressionAst variableExpression: + return variableExpression; + + // [type]$x = ... + case ConvertExpressionAst convertExpression: + typeConstraint = convertExpression.Type.TypeName.GetReflectionType(); + return GetVariableFromExpressionAst(convertExpression.Child, ref typeConstraint, ref setConstraint); + + // [attribute()][type]$x = ... + case AttributedExpressionAst attributedExpressionAst: + + try + { + setConstraint = attributedExpressionAst.Attribute.GetAttribute() as ValidateSetAttribute; + } + catch + { + // Do nothing, just prevent fallout from an unsuccessful attribute conversion + } + + return GetVariableFromExpressionAst(attributedExpressionAst.Child, ref typeConstraint, ref setConstraint); + + // Something else, like `MemberExpressionAst` $a.p = which isn't currently handled + default: + return null; + } + } + + // Gets any type constraints or validateset constraints on a given variable + private static bool TryGetTypeConstraintOnVariable( + CompletionContext completionContext, + string variableName, + out Type typeConstraint, + out ValidateSetAttribute setConstraint) + { + typeConstraint = null; + setConstraint = null; + + PSVariable variable = completionContext.ExecutionContext.EngineSessionState.GetVariable(variableName); + + if (variable == null || variable.Attributes.Count == 0) + { + return false; + } + + foreach (Attribute attribute in variable.Attributes) + { + if (attribute is ArgumentTypeConverterAttribute typeConverterAttribute) + { + typeConstraint = typeConverterAttribute.TargetType; + continue; + } + + if (attribute is ValidateSetAttribute validateSetAttribute) + { + setConstraint = validateSetAttribute; + } + } + + return typeConstraint != null || setConstraint != null; + } + + private static bool TryGetCompletionsForVariableAssignment( + CompletionContext completionContext, + AssignmentStatementAst assignmentAst, + out List completions) + { + bool TryGetResultForEnum(Type typeConstraint, CompletionContext completionContext, out List completions) + { + completions = null; + + if (typeConstraint != null && typeConstraint.IsEnum) + { + completions = GetResultForEnum(typeConstraint, completionContext); + return true; + } + + return false; + } + + bool TryGetResultForSet(Type typeConstraint, ValidateSetAttribute setConstraint, CompletionContext completionContext1, out List completions) + { + completions = null; + + if (setConstraint?.ValidValues != null) + { + completions = GetResultForSet(typeConstraint, setConstraint.ValidValues, completionContext); + return true; + } + + return false; + } + + completions = null; + + // Try to get the variable from the assignment, plus any type constraint on it + Type typeConstraint = null; + ValidateSetAttribute setConstraint = null; + VariableExpressionAst variableAst = GetVariableFromExpressionAst(assignmentAst.Left, ref typeConstraint, ref setConstraint); + + if (variableAst == null) + { + return false; + } + + // Assignment constraints override any existing ones, so try them first + + // Check any [ValidateSet()] constraint first since it's likely to be narrow + if (TryGetResultForSet(typeConstraint, setConstraint, completionContext, out completions)) + { + return true; + } + + // Then try to complete for an enum type + if (TryGetResultForEnum(typeConstraint, completionContext, out completions)) + { + return true; + } + + // If the assignment itself was unconstrained, the variable still might be + if (!TryGetTypeConstraintOnVariable(completionContext, variableAst.VariablePath.UserPath, out typeConstraint, out setConstraint)) + { + return false; + } + + // Again try the [ValidateSet()] constraint first + if (TryGetResultForSet(typeConstraint, setConstraint, completionContext, out completions)) + { + return true; + } + + // Then try to complete for an enum type again + if (TryGetResultForEnum(typeConstraint, completionContext, out completions)) + { + return true; + } + + return false; + } + + private static List GetResultForSet( + Type typeConstraint, + IList validValues, + CompletionContext completionContext) + { + var allValues = new List(); + foreach (string value in validValues) + { + if (typeConstraint != null && (typeConstraint == typeof(string) || typeConstraint.IsEnum)) + { + allValues.Add(GetQuotedString(value, completionContext)); + } + else + { + allValues.Add(value); + } + } + + return GetMatchedResults(allValues, completionContext); + } + + private static List GetMatchedResults( + List allValues, + CompletionContext completionContext) + { + var stringToComplete = string.Empty; + if (completionContext.TokenAtCursor != null && completionContext.TokenAtCursor.Kind != TokenKind.Equals) + { + stringToComplete = completionContext.TokenAtCursor.Text; + } + + IEnumerable matchedResults = null; + + if (!string.IsNullOrEmpty(stringToComplete)) + { + string matchString = stringToComplete + "*"; + var wildcardPattern = WildcardPattern.Get(matchString, WildcardOptions.IgnoreCase | WildcardOptions.CultureInvariant); + + matchedResults = allValues.Where(r => wildcardPattern.IsMatch(r)); + } + else + { + matchedResults = allValues; + } + + var result = new List(); + foreach (var match in matchedResults) + { + result.Add(new CompletionResult(match)); + } + + return result; + } + + private static string GetQuotedString( + string value, + CompletionContext completionContext) + { + var stringToComplete = string.Empty; + if (completionContext.TokenAtCursor != null) + { + stringToComplete = completionContext.TokenAtCursor.Text; + } + + var quote = stringToComplete.StartsWith('"') ? "\"" : "'"; + return quote + value + quote; + } + + private static List GetResultForEnum( + Type type, + CompletionContext completionContext) + { + var allNames = new List(); + foreach (var name in Enum.GetNames(type)) + { + allNames.Add(GetQuotedString(name, completionContext)); + } + + allNames.Sort(); + + return GetMatchedResults(allNames, completionContext); + } + private List GetResultForEnumPropertyValueOfDSCResource( CompletionContext completionContext, string stringToComplete, diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 index f1149f6b90..9f59011e6a 100644 --- a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 +++ b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 @@ -152,6 +152,144 @@ Describe "TabCompletion" -Tags CI { $res.CompletionMatches[1].CompletionText | Should -BeExactly '-Functionality' } + It 'Should work for variable assignment of enum type: ' -TestCases @( + @{ inputStr = '$ErrorActionPreference = '; filter = ''; doubleQuotes = $false } + @{ inputStr = '$ErrorActionPreference='; filter = ''; doubleQuotes = $false } + @{ inputStr = '$ErrorActionPreference="'; filter = ''; doubleQuotes = $true } + @{ inputStr = '$ErrorActionPreference = ''s'; filter = '| Where-Object { $_ -like "''s*" }'; doubleQuotes = $false } + @{ inputStr = '$ErrorActionPreference = "siL'; filter = '| Where-Object { $_ -like ''"sil*'' }'; doubleQuotes = $true } + @{ inputStr = '[System.Management.Automation.ActionPreference]$e='; filter = ''; doubleQuotes = $false } + @{ inputStr = '[System.Management.Automation.ActionPreference]$e = '; filter = ''; doubleQuotes = $false } + @{ inputStr = '[System.Management.Automation.ActionPreference]$e = "'; filter = ''; doubleQuotes = $true } + @{ inputStr = '[System.Management.Automation.ActionPreference]$e = "s'; filter = '| Where-Object { $_ -like """s*" }'; doubleQuotes = $true } + @{ inputStr = '[System.Management.Automation.ActionPreference]$e = "x'; filter = '| Where-Object { $_ -like """x*" }'; doubleQuotes = $true } + ){ + param($inputStr, $filter, $doubleQuotes) + + $quote = '''' + if ($doubleQuotes) { + $quote = '"' + } + + $sb = [scriptblock]::Create(@" + [cmdletbinding()] param([Parameter(ValueFromPipeline=`$true)]`$obj) process { `$obj $filter } +"@) + + $expectedValues = [enum]::GetValues("System.Management.Automation.ActionPreference") | ForEach-Object { $quote + $_.ToString() + $quote } | & $sb | Sort-Object + if ($expectedValues.Count -gt 0) { + $expected = [string]::Join(",",$expectedValues) + } + else { + $expected = '' + } + + $res = TabExpansion2 -inputScript $inputStr -cursorColumn $inputStr.Length + if ($res.CompletionMatches.Count -gt 0) { + $actual = [string]::Join(",",$res.CompletionMatches.completiontext) + } + else { + $actual = '' + } + + $actual | Should -BeExactly $expected + } + + It 'Should work for variable assignment of custom enum: ' -TestCases @( + @{ inputStr = '[Animal]$c="g'; expected = '"Giraffe"','"Goose"' } + @{ inputStr = '[Animal]$c='; expected = "'Duck'","'Giraffe'","'Goose'","'Horse'" } + @{ inputStr = '$script:test = "g'; expected = '"Giraffe"','"Goose"' } + @{ inputStr = '$script:test='; expected = "'Duck'","'Giraffe'","'Goose'","'Horse'" } + @{ inputStr = '$script:test = "x'; expected = @() } + ){ + param($inputStr, $expected) + + enum Animal { Duck; Goose; Horse; Giraffe } + [Animal]$script:test = 'Duck' + + $res = TabExpansion2 -inputScript $inputStr -cursorColumn $inputStr.Length + if ($res.CompletionMatches.Count -gt 0) { + $actual = [string]::Join(",",$res.CompletionMatches.completiontext) + } + else { + $actual = '' + } + + $actual | Should -BeExactly ([string]::Join(",",$expected)) + } + + It 'Should work for assignment of variable with validateset of strings: ' -TestCases @( + @{ inputStr = '$test='; expected = "'a'","'aa'","'aab'","'b'"; doubleQuotes = $false } + @{ inputStr = '$test="a'; expected = "'a'","'aa'","'aab'"; doubleQuotes = $true } + @{ inputStr = '$test = "aa'; expected = "'aa'","'aab'"; doubleQuotes = $true } + @{ inputStr = '$test=''aab'; expected = "'aab'"; doubleQuotes = $false } + @{ inputStr = '$test="c'; expected = ''; doubleQuotes = $true } + ){ + param($inputStr, $expected, $doubleQuotes) + + [ValidateSet('a','aa','aab','b')][string]$test = 'b' + + $expected = [string]::Join(",",$expected) + if ($doubleQuotes) { + $expected = $expected.Replace("'", """") + } + + $res = TabExpansion2 -inputScript $inputStr -cursorColumn $inputStr.Length + if ($res.CompletionMatches.Count -gt 0) { + $actual = [string]::Join(",",$res.CompletionMatches.completiontext) + } + else { + $actual = '' + } + + $actual | Should -BeExactly $expected + } + + It 'Should work for assignment of variable with validateset of int: ' -TestCases @( + @{ inputStr = '$test='; expected = 2,3,11,112 } + @{ inputStr = '$test = 1'; expected = 11,112 } + @{ inputStr = '$test =11'; expected = 11,112 } + @{ inputStr = '$test =4'; expected = @() } + ){ + param($inputStr, $expected) + + [ValidateSet(2,3,11,112)][int]$test = 2 + + $res = TabExpansion2 -inputScript $inputStr -cursorColumn $inputStr.Length + if ($res.CompletionMatches.Count -gt 0) { + $actual = [string]::Join(",",$res.CompletionMatches.completiontext) + } + else { + $actual = '' + } + + $actual | Should -BeExactly ([string]::Join(",",$expected)) + } + + It 'Should work for assignment of variable with validateset of strings: ' -TestCases @( + @{ inputStr = '[validateset("a","aa","aab","b")][string]$test='; expected = "'a'","'aa'","'aab'","'b'"; doubleQuotes = $false } + @{ inputStr = '[validateset("a","aa","aab","b")][string]$test="a'; expected = "'a'","'aa'","'aab'"; doubleQuotes = $true } + @{ inputStr = '[validateset("a","aa","aab","b")][string]$test = "aa'; expected = "'aa'","'aab'"; doubleQuotes = $true } + @{ inputStr = '[validateset("a","aa","aab","b")][string]$test=''aab'; expected = "'aab'"; doubleQuotes = $false } + @{ inputStr = '[validateset("a","aa","aab","b")][string]$test=''c'; expected = ''; doubleQuotes = $false } + ){ + param($inputStr, $expected, $doubleQuotes) + + $expected = [string]::Join(",",$expected) + if ($doubleQuotes) { + $expected = $expected.Replace("'", """") + } + + $res = TabExpansion2 -inputScript $inputStr -cursorColumn $inputStr.Length + if ($res.CompletionMatches.Count -gt 0) { + $actual = [string]::Join(",",$res.CompletionMatches.completiontext) + } + else { + $actual = '' + } + + $actual | Should -BeExactly $expected + } + Context NativeCommand { BeforeAll { $nativeCommand = (Get-Command -CommandType Application -TotalCount 1).Name