PSCompletion

1.0

PSCompletion.psm1

Add-Type @"
using System;
using System.Management.Automation;
 
public class NativeCommandTreeNode
{
    private NativeCommandTreeNode(NativeCommandTreeNode[] subCommands)
    {
        SubCommands = subCommands;
    }
 
    public NativeCommandTreeNode(string command, NativeCommandTreeNode[] subCommands)
        : this(command, null, subCommands)
    {
    }
 
    public NativeCommandTreeNode(string command, string tooltip, NativeCommandTreeNode[] subCommands)
        : this(subCommands)
    {
        this.Command = command;
        this.Tooltip = tooltip;
    }
 
    public NativeCommandTreeNode(string command, string tooltip, bool argument)
        : this(null)
    {
        this.Command = command;
        this.Tooltip = tooltip;
        this.Argument = true;
    }
 
    public NativeCommandTreeNode(ScriptBlock completionGenerator, NativeCommandTreeNode[] subCommands)
        : this(subCommands)
    {
        this.CompletionGenerator = completionGenerator;
    }
 
    public string Command { get; private set; }
    public string Tooltip { get; private set; }
    public bool Argument { get; private set; }
    public ScriptBlock CompletionGenerator { get; private set; }
    public NativeCommandTreeNode[] SubCommands { get; private set; }
}
"@

function New-CompletionResult
{
    param([Parameter(Position=0, ValueFromPipelineByPropertyName, Mandatory, ValueFromPipeline)]
          [ValidateNotNullOrEmpty()]
          [string]
          $CompletionText,

          [Parameter(Position=1, ValueFromPipelineByPropertyName)]
          [string]
          $ToolTip,

          [Parameter(Position=2, ValueFromPipelineByPropertyName)]
          [string]
          $ListItemText,

          [System.Management.Automation.CompletionResultType]
          $CompletionResultType = [System.Management.Automation.CompletionResultType]::Command,
 
          [Parameter(Mandatory = $false)]
          [switch] $NoQuotes = $false
          )

    process
    {
        $toolTipToUse = if ($ToolTip -eq '') { $CompletionText } else { $ToolTip }
        $listItemToUse = if ($ListItemText -eq '') { $CompletionText } else { $ListItemText }

        # If the caller explicitly requests that quotes
        # not be included, via the -NoQuotes parameter,
        # then skip adding quotes.

        if ($CompletionResultType -eq [System.Management.Automation.CompletionResultType]::Command -and -not $NoQuotes)
        {
            # Add single quotes for the caller in case they are needed.
            # We use the parser to robustly determine how it will treat
            # the argument. If we end up with too many tokens, or if
            # the parser found something expandable in the results, we
            # know quotes are needed.

            $tokens = $null
            $null = [System.Management.Automation.Language.Parser]::ParseInput("echo $CompletionText", [ref]$tokens, [ref]$null)
            if ($tokens.Length -ne 3 -or
                ($tokens[1] -is [System.Management.Automation.Language.StringExpandableToken] -and
                 $tokens[1].Kind -eq [System.Management.Automation.Language.TokenKind]::Generic))
            {
                $CompletionText = "'$CompletionText'"
            }
        }
        return New-Object System.Management.Automation.CompletionResult `
            ($CompletionText,$listItemToUse,$CompletionResultType,$toolTipToUse.Trim())
    }

}

function New-CommandTree
{
    [CmdletBinding(DefaultParameterSetName='Default')]
    param(
        [Parameter(Position=0, Mandatory, ParameterSetName='Default')]
        [Parameter(Position=0, Mandatory, ParameterSetName='Argument')]
        [ValidateNotNullOrEmpty()]
        [string]
        $Completion,

        [Parameter(Position=1, Mandatory, ParameterSetName='Default')]
        [Parameter(Position=1, Mandatory, ParameterSetName='Argument')]
        [string]
        $Tooltip,

        [Parameter(ParameterSetName='Argument')]
        [switch]
        $Argument,

        [Parameter(Position=2, ParameterSetName='Default')]
        [Parameter(Position=1, ParameterSetName='ScriptBlockSet')]
        [scriptblock]
        $SubCommands,

        [Parameter(Position=0, Mandatory, ParameterSetName='ScriptBlockSet')]
        [scriptblock]
        $CompletionGenerator
    )

    $actualSubCommands = $null
    if ($null -ne $SubCommands)
    {
        $actualSubCommands = [NativeCommandTreeNode[]](& $SubCommands)
    }

    switch ($PSCmdlet.ParameterSetName)
    {
        'Default' {
            New-Object NativeCommandTreeNode $Completion,$Tooltip,$actualSubCommands
            break
        }
        'Argument' {
            New-Object NativeCommandTreeNode $Completion,$Tooltip,$true
        }
        'ScriptBlockSet' {
            New-Object NativeCommandTreeNode $CompletionGenerator,$actualSubCommands
            break
        }
    }
}

function Get-CommandTreeCompletion
{
    param($wordToComplete, $commandAst, [NativeCommandTreeNode[]]$CommandTree)

    $commandElements = $commandAst.CommandElements

    # Skip the first command element - it's the command name
    # Iterate through the remaining elements, stopping early
    # if we find the element that matches $wordToComplete.
    for ($i = 1; $i -lt $commandElements.Count; $i++)
    {
        if (!($commandElements[$i] -is [System.Management.Automation.Language.StringConstantExpressionAst]))
        {
            # Ignore arguments that are expressions. In some rare cases this
            # could cause strange completions because the context is incorrect, e.g.:
            # $c = 'advfirewall'
            # netsh $c firewall
            # Here we would be in advfirewall firewall context, but we'd complete as
            # though we were in firewall context.
            continue
        }

        if ($commandElements[$i].Value -eq $wordToComplete)
        {
            $CommandTree = $CommandTree |
                Where-Object { $_.Command -like "$wordToComplete*" -or $_.CompletionGenerator -ne $null }
            break
        }

        foreach ($subCommand in $CommandTree)
        {
            if ($subCommand.Command -eq $commandElements[$i].Value)
            {
                if (!$subCommand.Argument)
                {
                    $CommandTree = $subCommand.SubCommands
                }
                break
            }
        }
    }

    if ($null -ne $CommandTree)
    {
        $CommandTree | ForEach-Object {
            if ($_.Command)
            {
                $toolTip = if ($_.Tooltip) { $_.Tooltip } else { $_.Command }
                $type = if ($_.Argument) { [System.Management.Automation.CompletionResultType]::ParameterValue } else { [System.Management.Automation.CompletionResultType]::Command }
                New-CompletionResult -CompletionText $_.Command -ToolTip $toolTip -CompletionResultType $type
            }
            else
            {
                & $_.CompletionGenerator $wordToComplete $commandAst
            }
        }
    }
}

New-Alias -Name 'nct' -Value 'New-CommandTree'
Export-ModuleMember -Function @('New-CommandTree', 'Get-CommandTreeCompletion') -Alias @('nct')