Skip to content

Commit 53805fb

Browse files
daxian-dbwdaxian-dbw
authored
Add the parameter Register-ArgumentCompleter -NativeFallback to support registering a cover-all completer for native commands (#25230)
1 parent 8d524ad commit 53805fb

File tree

5 files changed

+133
-33
lines changed

5 files changed

+133
-33
lines changed

‎src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2638,9 +2638,8 @@ private static ScriptBlock GetCustomArgumentCompleter(
26382638
}
26392639
}
26402640

2641-
var registeredCompleters = optionKey.Equals("NativeArgumentCompleters", StringComparison.OrdinalIgnoreCase)
2642-
? context.NativeArgumentCompleters
2643-
: context.CustomArgumentCompleters;
2641+
bool isNative = optionKey.Equals("NativeArgumentCompleters", StringComparison.OrdinalIgnoreCase);
2642+
var registeredCompleters = isNative ? context.NativeArgumentCompleters : context.CustomArgumentCompleters;
26442643

26452644
if (registeredCompleters != null)
26462645
{
@@ -2651,6 +2650,13 @@ private static ScriptBlock GetCustomArgumentCompleter(
26512650
return scriptBlock;
26522651
}
26532652
}
2653+
2654+
// For a native command, if a fallback completer is registered, then return it.
2655+
// For example, the 'Microsoft.PowerShell.UnixTabCompletion' module.
2656+
if (isNative && registeredCompleters.TryGetValue(RegisterArgumentCompleterCommand.FallbackCompleterKey, out scriptBlock))
2657+
{
2658+
return scriptBlock;
2659+
}
26542660
}
26552661

26562662
return null;

‎src/System.Management.Automation/engine/CommandCompletion/ExtensibleCompletion.cs‎

Lines changed: 70 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -172,66 +172,110 @@ public abstract class ArgumentCompleterFactoryAttribute : ArgumentCompleterAttri
172172
[Cmdlet(VerbsLifecycle.Register, "ArgumentCompleter", HelpUri = "https://go.microsoft.com/fwlink/?LinkId=528576")]
173173
public class RegisterArgumentCompleterCommand : PSCmdlet
174174
{
175+
private const string PowerShellSetName = "PowerShellSet";
176+
private const string NativeCommandSetName = "NativeCommandSet";
177+
private const string NativeFallbackSetName = "NativeFallbackSet";
178+
179+
// Use a key that is unlikely to be a file name or path to indicate the fallback completer for native commands.
180+
internal const string FallbackCompleterKey = "___ps::<native_fallback_key>@@___";
181+
175182
/// <summary>
183+
/// Gets or sets the command names for which the argument completer is registered.
176184
/// </summary>
177-
[Parameter(ParameterSetName = "NativeSet", Mandatory = true)]
178-
[Parameter(ParameterSetName = "PowerShellSet")]
185+
[Parameter(ParameterSetName = NativeCommandSetName, Mandatory = true)]
186+
[Parameter(ParameterSetName = PowerShellSetName)]
179187
[SuppressMessage("Microsoft.Performance", "CA1819:PropertiesShouldNotReturnArrays")]
180188
public string[] CommandName { get; set; }
181189

182190
/// <summary>
191+
/// Gets or sets the name of the parameter for which the argument completer is registered.
183192
/// </summary>
184-
[Parameter(ParameterSetName = "PowerShellSet", Mandatory = true)]
193+
[Parameter(ParameterSetName = PowerShellSetName, Mandatory = true)]
185194
public string ParameterName { get; set; }
186195

187196
/// <summary>
197+
/// Gets or sets the script block that will be executed to provide argument completions.
188198
/// </summary>
189199
[Parameter(Mandatory = true)]
190200
[AllowNull()]
191201
public ScriptBlock ScriptBlock { get; set; }
192202

193203
/// <summary>
204+
/// Indicates the argument completer is for native commands.
194205
/// </summary>
195-
[Parameter(ParameterSetName = "NativeSet")]
206+
[Parameter(ParameterSetName = NativeCommandSetName)]
196207
public SwitchParameter Native { get; set; }
197208

209+
/// <summary>
210+
/// Indicates the argument completer is a fallback for any native commands that don't have a completer registered.
211+
/// </summary>
212+
[Parameter(ParameterSetName = NativeFallbackSetName)]
213+
public SwitchParameter NativeFallback { get; set; }
214+
198215
/// <summary>
199216
/// </summary>
200217
protected override void EndProcessing()
201218
{
202219
Dictionary<string, ScriptBlock> completerDictionary;
203-
if (ParameterName != null)
220+
221+
if (ParameterSetName is NativeFallbackSetName)
204222
{
205-
completerDictionary = Context.CustomArgumentCompleters ??
206-
(Context.CustomArgumentCompleters = new Dictionary<string, ScriptBlock>(StringComparer.OrdinalIgnoreCase));
223+
completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase);
224+
225+
SetKeyValue(completerDictionary, FallbackCompleterKey, ScriptBlock);
207226
}
208-
else
227+
else if (ParameterSetName is NativeCommandSetName)
209228
{
210-
completerDictionary = Context.NativeArgumentCompleters ??
211-
(Context.NativeArgumentCompleters = new Dictionary<string, ScriptBlock>(StringComparer.OrdinalIgnoreCase));
212-
}
229+
completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase);
213230

214-
if (CommandName == null || CommandName.Length == 0)
231+
foreach (string command in CommandName)
232+
{
233+
var key = command?.Trim();
234+
if (string.IsNullOrEmpty(key))
235+
{
236+
continue;
237+
}
238+
239+
SetKeyValue(completerDictionary, key, ScriptBlock);
240+
}
241+
}
242+
else if (ParameterSetName is PowerShellSetName)
215243
{
216-
CommandName = new[] { string.Empty };
244+
completerDictionary = Context.CustomArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase);
245+
246+
string paramName = ParameterName.Trim();
247+
if (paramName.Length is 0)
248+
{
249+
return;
250+
}
251+
252+
if (CommandName is null || CommandName.Length is 0)
253+
{
254+
SetKeyValue(completerDictionary, paramName, ScriptBlock);
255+
return;
256+
}
257+
258+
foreach (string command in CommandName)
259+
{
260+
var key = command?.Trim();
261+
key = string.IsNullOrEmpty(key)
262+
? paramName
263+
: $"{key}:{paramName}";
264+
265+
SetKeyValue(completerDictionary, key, ScriptBlock);
266+
}
217267
}
218268

219-
for (int i = 0; i < CommandName.Length; i++)
269+
static void SetKeyValue(Dictionary<string, ScriptBlock> table, string key, ScriptBlock value)
220270
{
221-
var key = CommandName[i];
222-
if (!string.IsNullOrWhiteSpace(ParameterName))
271+
if (value is null)
223272
{
224-
if (!string.IsNullOrWhiteSpace(key))
225-
{
226-
key = key + ":" + ParameterName;
227-
}
228-
else
229-
{
230-
key = ParameterName;
231-
}
273+
table.Remove(key);
274+
}
275+
else
276+
{
277+
table[key] = value;
232278
}
233-
234-
completerDictionary[key] = ScriptBlock;
235279
}
236280
}
237281
}

‎src/System.Management.Automation/engine/hostifaces/LocalPipeline.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1067,9 +1067,9 @@ internal override void SetHistoryString(string historyString)
10671067
/// ExecutionContext, if it available in TLS
10681068
/// Null, if ExecutionContext is not available in TLS
10691069
/// </returns>
1070-
internal static System.Management.Automation.ExecutionContext GetExecutionContextFromTLS()
1070+
internal static ExecutionContext GetExecutionContextFromTLS()
10711071
{
1072-
System.Management.Automation.Runspaces.Runspace runspace = Runspace.DefaultRunspace;
1072+
Runspace runspace = Runspace.DefaultRunspace;
10731073
if (runspace == null)
10741074
{
10751075
return null;

‎src/System.Management.Automation/engine/lang/scriptblock.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ internal static ScriptBlock Create(ExecutionContext context, string script)
104104
/// </summary>
105105
/// <param name="script">The string to compile.</param>
106106
public static ScriptBlock Create(string script) => Create(
107-
parser: new Language.Parser(),
107+
parser: new Parser(),
108108
fileName: null,
109109
fileContents: script);
110110

‎test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1929,8 +1929,12 @@ param([ValidatePattern(
19291929

19301930
Context NativeCommand {
19311931
BeforeAll {
1932-
$nativeCommand = (Get-Command -CommandType Application -TotalCount 1).Name
1932+
## Find a native command that is not 'pwsh'. We will use 'pwsh' for fallback completer tests later.
1933+
$nativeCommand = Get-Command -CommandType Application -TotalCount 2 |
1934+
Where-Object Name -NotLike pwsh* |
1935+
Select-Object -First 1
19331936
}
1937+
19341938
It 'Completes native commands with -' {
19351939
Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock {
19361940
param($wordToComplete, $ast, $cursorColumn)
@@ -1994,6 +1998,52 @@ param([ValidatePattern(
19941998
$res.CompletionMatches | Should -HaveCount 1
19951999
$res.CompletionMatches.CompletionText | Should -BeExactly "-option"
19962000
}
2001+
2002+
It 'Covers an arbitrary unbound native command with -t' {
2003+
## Register a completer for $nativeCommand.
2004+
Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock {
2005+
param($wordToComplete, $ast, $cursorColumn)
2006+
if ($wordToComplete -eq '-t') {
2007+
return "-terminal"
2008+
}
2009+
}
2010+
2011+
## Register a fallback native command completer.
2012+
Register-ArgumentCompleter -NativeFallback -ScriptBlock {
2013+
param($wordToComplete, $ast, $cursorColumn)
2014+
if ($wordToComplete -eq '-t') {
2015+
return "-testing"
2016+
}
2017+
}
2018+
2019+
## The specific completer will be used if it exists.
2020+
$line = "$nativeCommand -t"
2021+
$res = TabExpansion2 -inputScript $line -cursorColumn $line.Length
2022+
$res.CompletionMatches | Should -HaveCount 1
2023+
$res.CompletionMatches.CompletionText | Should -BeExactly "-terminal"
2024+
2025+
## Otherwise, the fallback completer will kick in.
2026+
$line = "pwsh -t"
2027+
$res = TabExpansion2 -inputScript $line -cursorColumn $line.Length
2028+
$res.CompletionMatches | Should -HaveCount 1
2029+
$res.CompletionMatches.CompletionText | Should -BeExactly "-testing"
2030+
2031+
## Remove the completer for $nativeCommand.
2032+
Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock $null
2033+
2034+
## The fallback completer will be used for $nativeCommand.
2035+
$line = "$nativeCommand -t"
2036+
$res = TabExpansion2 -inputScript $line -cursorColumn $line.Length
2037+
$res.CompletionMatches | Should -HaveCount 1
2038+
$res.CompletionMatches.CompletionText | Should -BeExactly "-testing"
2039+
2040+
## Remove the fallback completer for $nativeCommand.
2041+
Register-ArgumentCompleter -NativeFallback -ScriptBlock $null
2042+
2043+
## The fallback completer will be used for $nativeCommand.
2044+
$res = TabExpansion2 -inputScript $line -cursorColumn $line.Length
2045+
$res.CompletionMatches | Should -HaveCount 0
2046+
}
19972047
}
19982048

19992049
It 'Should complete "Export-Counter -FileFormat" with available output formats' -Pending {

0 commit comments

Comments
 (0)