MegaBites-AI/Windows-powershell

1// Copyright (c) Microsoft Corporation.2// Licensed under the MIT License.3 4using System.Collections;5using System.Collections.Generic;6using System.Diagnostics.CodeAnalysis;7using System.Linq;8using System.Management.Automation.Language;9 10namespace System.Management.Automation11{12    /// <summary>13    /// This attribute is used to specify an argument completer for a parameter to a cmdlet or function.14    /// <example>15    /// <code>16    ///     [Parameter()]17    ///     [ArgumentCompleter(typeof(NounArgumentCompleter))]18    ///     public string Noun { get; set; }19    /// </code>20    /// </example>21    /// </summary>22    [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property)]23    public class ArgumentCompleterAttribute : Attribute24    {25        /// <summary/>26        [SuppressMessage("Microsoft.Naming", "CA1721:PropertyNamesShouldNotMatchGetMethods")]27        public Type Type { get; }28 29        /// <summary/>30        public ScriptBlock ScriptBlock { get; }31 32        /// <param name="type">The type must implement <see cref="IArgumentCompleter"/> and have a default constructor.</param>33        public ArgumentCompleterAttribute(Type type)34        {35            if (type == null || (type.GetInterfaces().All(static t => t != typeof(IArgumentCompleter))))36            {37                throw PSTraceSource.NewArgumentException(nameof(type));38            }39 40            Type = type;41        }42 43        /// <summary>44        /// Initializes a new instance of the <see cref="ArgumentCompleterAttribute"/> class.45        /// This constructor is used by derived attributes implementing <see cref="IArgumentCompleterFactory"/>.46        /// </summary>47        protected ArgumentCompleterAttribute()48        {49            if (this is not IArgumentCompleterFactory)50            {51                throw PSTraceSource.NewInvalidOperationException();52            }53        }54 55        /// <summary>56        /// This constructor is used primarily via PowerShell scripts.57        /// </summary>58        /// <param name="scriptBlock"></param>59        public ArgumentCompleterAttribute(ScriptBlock scriptBlock)60        {61            if (scriptBlock is null)62            {63                throw PSTraceSource.NewArgumentNullException(nameof(scriptBlock));64            }65 66            ScriptBlock = scriptBlock;67        }68 69        internal IArgumentCompleter CreateArgumentCompleter()70        {71            return Type != null72                ? Activator.CreateInstance(Type) as IArgumentCompleter73                : this is IArgumentCompleterFactory factory74                    ? factory.Create()75                    : null;76        }77    }78 79    /// <summary>80    /// A type specified by the <see cref="ArgumentCompleterAttribute"/> must implement this interface.81    /// </summary>82#nullable enable83    public interface IArgumentCompleter84    {85        /// <summary>86        /// Implementations of this function are called by PowerShell to complete arguments.87        /// </summary>88        /// <param name="commandName">The name of the command that needs argument completion.</param>89        /// <param name="parameterName">The name of the parameter that needs argument completion.</param>90        /// <param name="wordToComplete">The (possibly empty) word being completed.</param>91        /// <param name="commandAst">The command ast in case it is needed for completion.</param>92        /// <param name="fakeBoundParameters">93        /// This parameter is similar to $PSBoundParameters, except that sometimes PowerShell cannot or94        /// will not attempt to evaluate an argument, in which case you may need to use <paramref name="commandAst"/>.95        /// </param>96        /// <returns>97        /// A collection of completion results, most like with <see cref="CompletionResult.ResultType"/> set to98        /// <see cref="CompletionResultType.ParameterValue"/>.99        /// </returns>100        IEnumerable<CompletionResult> CompleteArgument(101            string commandName,102            string parameterName,103            string wordToComplete,104            CommandAst commandAst,105            IDictionary fakeBoundParameters);106    }107#nullable restore108 109    /// <summary>110    /// Creates a new argument completer.111    /// </summary>112    /// <para>113    /// If an attribute that derives from <see cref="ArgumentCompleterAttribute"/> implements this interface,114    /// it will be used to create the <see cref="IArgumentCompleter"/>, thus giving a way to parameterize a completer.115    /// The derived attribute can have properties or constructor arguments that are used when creating the completer.116    /// </para>117    /// <example>118    /// This example shows the intended usage of <see cref="IArgumentCompleterFactory"/> to pass arguments to an argument completer.119    /// <code>120    /// public class NumberCompleterAttribute : ArgumentCompleterAttribute, IArgumentCompleterFactory {121    ///    private readonly int _from;122    ///    private readonly int _to;123    ///124    ///    public NumberCompleterAttribute(int from, int to){125    ///       _from = from;126    ///       _to = to;127    ///    }128    ///129    ///    // use the attribute parameters to create a parameterized completer130    ///    IArgumentCompleter Create() => new NumberCompleter(_from, _to);131    /// }132    ///133    /// class NumberCompleter : IArgumentCompleter {134    ///    private readonly int _from;135    ///    private readonly int _to;136    ///137    ///    public NumberCompleter(int from, int to){138    ///       _from = from;139    ///       _to = to;140    ///    }141    ///142    ///    IEnumerable{CompletionResult} CompleteArgument(string commandName, string parameterName, string wordToComplete,143    ///       CommandAst commandAst, IDictionary fakeBoundParameters) {144    ///       for(int i = _from; i &lt; _to; i++) {145    ///           yield return new CompletionResult(i.ToString());146    ///       }147    ///    }148    /// }149    /// </code>150    /// </example>151    public interface IArgumentCompleterFactory152    {153        /// <summary>154        /// Creates an instance of a class implementing the <see cref="IArgumentCompleter"/> interface.155        /// </summary>156        /// <returns>An IArgumentCompleter instance.</returns>157        IArgumentCompleter Create();158    }159 160    /// <summary>161    /// Base class for parameterized argument completer attributes.162    /// </summary>163    [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property)]164    public abstract class ArgumentCompleterFactoryAttribute : ArgumentCompleterAttribute, IArgumentCompleterFactory165    {166        /// <inheritdoc />167        public abstract IArgumentCompleter Create();168    }169 170    /// <summary>171    /// </summary>172    [Cmdlet(VerbsLifecycle.Register, "ArgumentCompleter", HelpUri = "https://go.microsoft.com/fwlink/?LinkId=528576")]173    public class RegisterArgumentCompleterCommand : PSCmdlet174    {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 182        /// <summary>183        /// Gets or sets the command names for which the argument completer is registered.184        /// </summary>185        [Parameter(ParameterSetName = NativeCommandSetName, Mandatory = true)]186        [Parameter(ParameterSetName = PowerShellSetName)]187        [SuppressMessage("Microsoft.Performance", "CA1819:PropertiesShouldNotReturnArrays")]188        public string[] CommandName { get; set; }189 190        /// <summary>191        /// Gets or sets the name of the parameter for which the argument completer is registered.192        /// </summary>193        [Parameter(ParameterSetName = PowerShellSetName, Mandatory = true)]194        public string ParameterName { get; set; }195 196        /// <summary>197        /// Gets or sets the script block that will be executed to provide argument completions.198        /// </summary>199        [Parameter(Mandatory = true)]200        [AllowNull()]201        public ScriptBlock ScriptBlock { get; set; }202 203        /// <summary>204        /// Indicates the argument completer is for native commands.205        /// </summary>206        [Parameter(ParameterSetName = NativeCommandSetName)]207        public SwitchParameter Native { get; set; }208 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 215        /// <summary>216        /// </summary>217        protected override void EndProcessing()218        {219            Dictionary<string, ScriptBlock> completerDictionary;220 221            if (ParameterSetName is NativeFallbackSetName)222            {223                completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase);224 225                SetKeyValue(completerDictionary, FallbackCompleterKey, ScriptBlock);226            }227            else if (ParameterSetName is NativeCommandSetName)228            {229                completerDictionary = Context.NativeArgumentCompleters ??= new(StringComparer.OrdinalIgnoreCase);230 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)243            {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                        ? paramName263                        : $"{key}:{paramName}";264 265                    SetKeyValue(completerDictionary, key, ScriptBlock);266                }267            }268 269            static void SetKeyValue(Dictionary<string, ScriptBlock> table, string key, ScriptBlock value)270            {271                if (value is null)272                {273                    table.Remove(key);274                }275                else276                {277                    table[key] = value;278                }279            }280        }281    }282 283    /// <summary>284    /// This attribute is used to specify an argument completions for a parameter of a cmdlet or function285    /// based on string array.286    /// <example>287    ///     [Parameter()]288    ///     [ArgumentCompletions("Option1","Option2","Option3")]289    ///     public string Noun { get; set; }290    /// </example>291    /// </summary>292    [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property)]293    public class ArgumentCompletionsAttribute : Attribute294    {295        private readonly string[] _completions;296 297        /// <summary>298        /// Initializes a new instance of the ArgumentCompletionsAttribute class.299        /// </summary>300        /// <param name="completions">List of complete values.</param>301        /// <exception cref="ArgumentNullException">For null arguments.</exception>302        /// <exception cref="ArgumentOutOfRangeException">For invalid arguments.</exception>303        public ArgumentCompletionsAttribute(params string[] completions)304        {305            if (completions == null)306            {307                throw PSTraceSource.NewArgumentNullException(nameof(completions));308            }309 310            if (completions.Length == 0)311            {312                throw PSTraceSource.NewArgumentOutOfRangeException(nameof(completions), completions);313            }314 315            _completions = completions;316        }317 318        /// <summary>319        /// The function returns completions for arguments.320        /// </summary>321        public IEnumerable<CompletionResult> CompleteArgument(string commandName, string parameterName, string wordToComplete, CommandAst commandAst, IDictionary fakeBoundParameters)322        {323            var wordToCompletePattern = WildcardPattern.Get(string.IsNullOrWhiteSpace(wordToComplete) ? "*" : wordToComplete + "*", WildcardOptions.IgnoreCase);324 325            foreach (var str in _completions)326            {327                if (wordToCompletePattern.IsMatch(str))328                {329                    yield return new CompletionResult(str, str, CompletionResultType.ParameterValue, str);330                }331            }332        }333    }334}335