From 53805fb426d4e60601c9013287cbc2690aaef1e0 Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Fri, 29 Aug 2025 15:50:02 -0700 Subject: [PATCH] Add the parameter `Register-ArgumentCompleter -NativeFallback` to support registering a cover-all completer for native commands (#25230) --- .../CommandCompletion/CompletionCompleters.cs | 12 ++- .../CommandCompletion/ExtensibleCompletion.cs | 98 ++++++++++++++----- .../engine/hostifaces/LocalPipeline.cs | 4 +- .../engine/lang/scriptblock.cs | 2 +- .../TabCompletion/TabCompletion.Tests.ps1 | 52 +++++++++- 5 files changed, 134 insertions(+), 34 deletions(-) diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs index 880aedb04c..10ce5c851e 100644 --- a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs +++ b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs @@ -2638,9 +2638,8 @@ namespace System.Management.Automation } } - var registeredCompleters = optionKey.Equals("NativeArgumentCompleters", StringComparison.OrdinalIgnoreCase) - ? context.NativeArgumentCompleters - : context.CustomArgumentCompleters; + bool isNative = optionKey.Equals("NativeArgumentCompleters", StringComparison.OrdinalIgnoreCase); + var registeredCompleters = isNative ? context.NativeArgumentCompleters : context.CustomArgumentCompleters; if (registeredCompleters != null) { @@ -2651,6 +2650,13 @@ namespace System.Management.Automation return scriptBlock; } } + + // For a native command, if a fallback completer is registered, then return it. + // For example, the 'Microsoft.PowerShell.UnixTabCompletion' module. + if (isNative && registeredCompleters.TryGetValue(RegisterArgumentCompleterCommand.FallbackCompleterKey, out scriptBlock)) + { + return scriptBlock; + } } return null; diff --git a/src/System.Management.Automation/engine/CommandCompletion/ExtensibleCompletion.cs b/src/System.Management.Automation/engine/CommandCompletion/ExtensibleCompletion.cs index 2b5281a8f2..c63a8e7f92 100644 --- a/src/System.Management.Automation/engine/CommandCompletion/ExtensibleCompletion.cs +++ b/src/System.Management.Automation/engine/CommandCompletion/ExtensibleCompletion.cs @@ -172,66 +172,110 @@ namespace System.Management.Automation [Cmdlet(VerbsLifecycle.Register, "ArgumentCompleter", HelpUri = "https://go.microsoft.com/fwlink/?LinkId=528576")] public class RegisterArgumentCompleterCommand : PSCmdlet { + private const string PowerShellSetName = "PowerShellSet"; + private const string NativeCommandSetName = "NativeCommandSet"; + private const string NativeFallbackSetName = "NativeFallbackSet"; + + // Use a key that is unlikely to be a file name or path to indicate the fallback completer for native commands. + internal const string FallbackCompleterKey = "___ps::@@___"; + /// + /// Gets or sets the command names for which the argument completer is registered. /// - [Parameter(ParameterSetName = "NativeSet", Mandatory = true)] - [Parameter(ParameterSetName = "PowerShellSet")] + [Parameter(ParameterSetName = NativeCommandSetName, Mandatory = true)] + [Parameter(ParameterSetName = PowerShellSetName)] [SuppressMessage("Microsoft.Performance", "CA1819:PropertiesShouldNotReturnArrays")] public string[] CommandName { get; set; } /// + /// Gets or sets the name of the parameter for which the argument completer is registered. /// - [Parameter(ParameterSetName = "PowerShellSet", Mandatory = true)] + [Parameter(ParameterSetName = PowerShellSetName, Mandatory = true)] public string ParameterName { get; set; } /// + /// Gets or sets the script block that will be executed to provide argument completions. /// [Parameter(Mandatory = true)] [AllowNull()] public ScriptBlock ScriptBlock { get; set; } /// + /// Indicates the argument completer is for native commands. /// - [Parameter(ParameterSetName = "NativeSet")] + [Parameter(ParameterSetName = NativeCommandSetName)] public SwitchParameter Native { get; set; } + /// + /// Indicates the argument completer is a fallback for any native commands that don't have a completer registered. + /// + [Parameter(ParameterSetName = NativeFallbackSetName)] + public SwitchParameter NativeFallback { get; set; } + /// /// protected override void EndProcessing() { Dictionary completerDictionary; - if (ParameterName != null) - { - completerDictionary = Context.CustomArgumentCompleters ?? - (Context.CustomArgumentCompleters = new Dictionary(StringComparer.OrdinalIgnoreCase)); - } - else - { - completerDictionary = Context.NativeArgumentCompleters ?? - (Context.NativeArgumentCompleters = new Dictionary(StringComparer.OrdinalIgnoreCase)); - } - if (CommandName == null || CommandName.Length == 0) + if (ParameterSetName is NativeFallbackSetName) { - CommandName = new[] { string.Empty }; - } + completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase); - for (int i = 0; i < CommandName.Length; i++) + SetKeyValue(completerDictionary, FallbackCompleterKey, ScriptBlock); + } + else if (ParameterSetName is NativeCommandSetName) { - var key = CommandName[i]; - if (!string.IsNullOrWhiteSpace(ParameterName)) + completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase); + + foreach (string command in CommandName) { - if (!string.IsNullOrWhiteSpace(key)) + var key = command?.Trim(); + if (string.IsNullOrEmpty(key)) { - key = key + ":" + ParameterName; - } - else - { - key = ParameterName; + continue; } + + SetKeyValue(completerDictionary, key, ScriptBlock); + } + } + else if (ParameterSetName is PowerShellSetName) + { + completerDictionary = Context.CustomArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase); + + string paramName = ParameterName.Trim(); + if (paramName.Length is 0) + { + return; } - completerDictionary[key] = ScriptBlock; + if (CommandName is null || CommandName.Length is 0) + { + SetKeyValue(completerDictionary, paramName, ScriptBlock); + return; + } + + foreach (string command in CommandName) + { + var key = command?.Trim(); + key = string.IsNullOrEmpty(key) + ? paramName + : $"{key}:{paramName}"; + + SetKeyValue(completerDictionary, key, ScriptBlock); + } + } + + static void SetKeyValue(Dictionary table, string key, ScriptBlock value) + { + if (value is null) + { + table.Remove(key); + } + else + { + table[key] = value; + } } } } diff --git a/src/System.Management.Automation/engine/hostifaces/LocalPipeline.cs b/src/System.Management.Automation/engine/hostifaces/LocalPipeline.cs index 93cb65f3f0..a9a3a66b85 100644 --- a/src/System.Management.Automation/engine/hostifaces/LocalPipeline.cs +++ b/src/System.Management.Automation/engine/hostifaces/LocalPipeline.cs @@ -1067,9 +1067,9 @@ namespace System.Management.Automation.Runspaces /// ExecutionContext, if it available in TLS /// Null, if ExecutionContext is not available in TLS /// - internal static System.Management.Automation.ExecutionContext GetExecutionContextFromTLS() + internal static ExecutionContext GetExecutionContextFromTLS() { - System.Management.Automation.Runspaces.Runspace runspace = Runspace.DefaultRunspace; + Runspace runspace = Runspace.DefaultRunspace; if (runspace == null) { return null; diff --git a/src/System.Management.Automation/engine/lang/scriptblock.cs b/src/System.Management.Automation/engine/lang/scriptblock.cs index 7bbd33c9cf..06fa66d0f7 100644 --- a/src/System.Management.Automation/engine/lang/scriptblock.cs +++ b/src/System.Management.Automation/engine/lang/scriptblock.cs @@ -104,7 +104,7 @@ namespace System.Management.Automation /// /// The string to compile. public static ScriptBlock Create(string script) => Create( - parser: new Language.Parser(), + parser: new Parser(), fileName: null, fileContents: script); diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 index 8e4fdd9d66..046d00b442 100644 --- a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 +++ b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 @@ -1929,8 +1929,12 @@ param([ValidatePattern( Context NativeCommand { BeforeAll { - $nativeCommand = (Get-Command -CommandType Application -TotalCount 1).Name + ## Find a native command that is not 'pwsh'. We will use 'pwsh' for fallback completer tests later. + $nativeCommand = Get-Command -CommandType Application -TotalCount 2 | + Where-Object Name -NotLike pwsh* | + Select-Object -First 1 } + It 'Completes native commands with -' { Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock { param($wordToComplete, $ast, $cursorColumn) @@ -1994,6 +1998,52 @@ param([ValidatePattern( $res.CompletionMatches | Should -HaveCount 1 $res.CompletionMatches.CompletionText | Should -BeExactly "-option" } + + It 'Covers an arbitrary unbound native command with -t' { + ## Register a completer for $nativeCommand. + Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock { + param($wordToComplete, $ast, $cursorColumn) + if ($wordToComplete -eq '-t') { + return "-terminal" + } + } + + ## Register a fallback native command completer. + Register-ArgumentCompleter -NativeFallback -ScriptBlock { + param($wordToComplete, $ast, $cursorColumn) + if ($wordToComplete -eq '-t') { + return "-testing" + } + } + + ## The specific completer will be used if it exists. + $line = "$nativeCommand -t" + $res = TabExpansion2 -inputScript $line -cursorColumn $line.Length + $res.CompletionMatches | Should -HaveCount 1 + $res.CompletionMatches.CompletionText | Should -BeExactly "-terminal" + + ## Otherwise, the fallback completer will kick in. + $line = "pwsh -t" + $res = TabExpansion2 -inputScript $line -cursorColumn $line.Length + $res.CompletionMatches | Should -HaveCount 1 + $res.CompletionMatches.CompletionText | Should -BeExactly "-testing" + + ## Remove the completer for $nativeCommand. + Register-ArgumentCompleter -Native -CommandName $nativeCommand -ScriptBlock $null + + ## The fallback completer will be used for $nativeCommand. + $line = "$nativeCommand -t" + $res = TabExpansion2 -inputScript $line -cursorColumn $line.Length + $res.CompletionMatches | Should -HaveCount 1 + $res.CompletionMatches.CompletionText | Should -BeExactly "-testing" + + ## Remove the fallback completer for $nativeCommand. + Register-ArgumentCompleter -NativeFallback -ScriptBlock $null + + ## The fallback completer will be used for $nativeCommand. + $res = TabExpansion2 -inputScript $line -cursorColumn $line.Length + $res.CompletionMatches | Should -HaveCount 0 + } } It 'Should complete "Export-Counter -FileFormat" with available output formats' -Pending {