NativeCommandCompletion/NativeCommandCompletion.psm1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class NativeCommandCompletionResult : System.Management.Automation.CompletionResult {
    [string]$TextType

    NativeCommandCompletionResult(
        [string]$CompletionText,
        [string]$ListItemText,
        [System.Management.Automation.CompletionResultType]$ResultType,
        [string]$ToolTip,
        [string]$TextType) :
    base($CompletionText, $ListItemText, $ResultType, $ToolTip) {

        $this.TextType = $TextType
    }
}

$nativeCommandCompleters = @{}

function New-CompletionResult {
    Param(
        [Parameter(Mandatory = $true)]
        [string]$CompletionText,
        [string]$TextType,
        [string]$ToolTip = $CompletionText,
        [string]$ListItemText = $CompletionText,
        [System.Management.Automation.CompletionResultType]$ResultType
    )

    if (!$ResultType) {
        if ($CompletionText.StartsWith('-')) {
            $ResultType = [System.Management.Automation.CompletionResultType]::ParameterName
        } elseif ($TextType -like '*Command') {
            $ResultType = [System.Management.Automation.CompletionResultType]::Command
        } else {
            $ResultType = [System.Management.Automation.CompletionResultType]::Text
        }
    }

    New-Object NativeCommandCompletionResult $CompletionText, $ListItemText, $ResultType, $ToolTip, $TextType
}

Set-Alias -Name COMPGEN -Value New-CompletionResult

function Register-Completer {
    Param(
        [Parameter(Mandatory = $true)]
        [string]$Name,
        $Completer,
        [Alias('Option')]
        $Parameter
    )

    if (!$nativeCommandCompleters.ContainsKey($Name)) {
        $nativeCommandCompleters[$Name] = @{}
    }

    if ($Completer) {
        $nativeCommandCompleters[$Name].completer = $Completer
    }
    if ($Parameter) {
        $nativeCommandCompleters[$Name].parameter = $Parameter
    }
}

function Get-Completer {
    Param(
        [Parameter(Mandatory = $true)]
        [string]$Name,
        [Alias('Option')]
        [switch]$Parameter
    )

    if ($Parameter) {
        $nativeCommandCompleters[$Name].parameter
    } else {
        $nativeCommandCompleters[$Name].completer
    }
}

function Invoke-Completer {
    Param(
        [Parameter(Mandatory = $true, Position = 0)]
        [string]$Name,
        [Parameter(Mandatory = $true, ParameterSetName = 'Completer')]
        [switch]$Completer,
        [Parameter(Mandatory = $true, ParameterSetName = 'Parameter')]
        [Alias('Option')]
        [switch]$Parameter,
        [Object[]]$ArgumentList
    )

    $cmpltr = switch ($PSCmdlet.ParameterSetName) {
        Completer { $nativeCommandCompleters[$Name].completer }
        Parameter { $nativeCommandCompleters[$Name].parameter }
    }

    if ($cmpltr -is [scriptblock]) {
        $cmpltr = Invoke-Command -ScriptBlock $cmpltr -ArgumentList $ArgumentList
    }

    foreach ($c in $cmpltr) {
        if ($c -is [NativeCommandCompletionResult]) {
            $c
        } else {
            New-CompletionResult -CompletionText $c -TextType Text -ResultType ([System.Management.Automation.CompletionResultType]::Text)
        }
    }
}

function Register-NativeCommandArgumentCompleter {
    Param(
        [Parameter(Mandatory = $true)]
        [string]$CommandName,
        [Parameter(Mandatory = $true)]
        [scriptblock]$ScriptBlock
    )

    $CommandName, (Get-Alias -Definition $CommandName -ErrorAction SilentlyContinue).Name | ForEach-Object {
        if ($_) {
            Register-ArgumentCompleter -CommandName $_ -ScriptBlock $ScriptBlock -Native
        }
    }
}