From 16927905069f2a5779c388c801f3af2f7e315486 Mon Sep 17 00:00:00 2001
From: Evan Fiordeliso <evan.fiordeliso@gmail.com>
Date: Mon, 13 Nov 2023 00:44:09 -0500
Subject: [PATCH] Better fish completion using functions

---
 completions_fish.go | 81 ++++++++++++++++++++++++++++++---------------
 tpl_funcs.go        | 18 ++++++++--
 2 files changed, 71 insertions(+), 28 deletions(-)

diff --git a/completions_fish.go b/completions_fish.go
index 09eff7a..c01d608 100644
--- a/completions_fish.go
+++ b/completions_fish.go
@@ -15,59 +15,88 @@ func WriteFishCompletions(out io.Writer, rootCmd *Command) error {
 }
 
 var fishTpl = template.Must(template.New("fish").Funcs(tplFuncs).Parse(`
-set -l progName {{ .RootCmd.Name }}
-set -l commands {{- range .RootCmd.Subcommands }} {{ .Name }}{{ end }}
+{{- $rootCmd := .RootCmd -}}
+{{- $progName := $rootCmd.Name -}}
+{{- $varName := under $rootCmd.Name -}}
+set -l commands {{- range $rootCmd.Subcommands }} {{ .Name }}{{ end }}
 
-{{- /* Option template */ -}}
+complete -c {{ $progName }} -f
+
+function __fish_{{ $varName }}_needs_command
+	set -l cmd (commandline -opc)
+	if test (count $cmd) -eq 1
+		return 0
+	end
+	return 1
+end
+
+function __fish_{{ $varName }}_using_command
+	set -l cmd (commandline -opc)
+	echo $cmd
+	set -l cnt (count $argv)
+	if test (count $cmd) -gt $cnt
+		for i in (seq 1 $cnt);
+			if test $argv[$i] != $cmd[(math $i + 1)]
+				return 1
+			end
+		end
+		return 0
+	end
+	return 1
+end
+
+{{/* Option template */ -}}
 {{ define "opt" }}
-complete -c $progName
+{{ $progName := .ProgName -}}
+{{- $varName := .VarName -}}
+complete -c {{ $progName }}
 {{- if ne .Opt.ShortName "" }} -s {{ .Opt.ShortName }} {{ end -}}
 {{- if ne .Opt.Name "" }} -l {{ .Opt.Name }} {{ end -}} 
-{{- if .Cond }} -n "{{ .Cond }}" {{ end -}}
+{{- if .Cmd }} {{- if ne .Cmd.CommandPath "" -}}
+-n "__fish_{{ $varName }}_using_command {{.Cmd.CommandPath}}" {{ else -}}
+-n "__fish_{{ $varName }}_needs_command" {{ end -}}{{- end -}}
 -d '{{ .Opt.Description }}'
 {{- end }}
 
 {{- /* Command template */ -}}
 {{ define "cmd" }}
-{{ $parentVarPrefix := "" -}}
-{{- $varPrefix := join .Cmd.Name "_" -}}
-{{- if .VarPrefix -}}
-{{- $varPrefix = join .VarPrefix $varPrefix -}}
-{{- $parentVarPrefix = .VarPrefix -}}
-{{- end -}}
+{{ $progName := .ProgName -}}
+{{- $varName := .VarName -}}
+{{- $cmd := .Cmd -}}
+{{- $parentVarPrefix := varPrefix .Cmd.Parent.CommandPath -}}
+{{- $varPrefix := varPrefix .Cmd.CommandPath -}}
 
-{{ $parentCond := "" }}
-{{- $cond := join "__fish_seen_subcommand_from " .Cmd.Name }}
-{{- if .Prefix -}}
-{{- $cond = join .Prefix $cond -}}
-{{- $parentCond = .Prefix -}}
+{{- if eq .Cmd.Parent.CommandPath "" -}}
+complete -f -c {{ $progName }} -n "__fish_{{ $varName }}_needs_command" -a {{ .Cmd.Name }} -d "{{ .Cmd.ShortDescription }}"
+{{- else -}}
+complete -f -c {{ $progName }} -n "__fish_{{ $varName }}_using_command {{.Cmd.Parent.CommandPath}}" -a {{ .Cmd.Name }} -d "{{ .Cmd.ShortDescription }}"
 {{- end -}}
 
-set -l {{ $varPrefix }}commands {{- range .Cmd.Subcommands }} {{ .Name }}{{ end }}
-complete -f -c $progName -n "{{ $parentCond }}not __fish_seen_subcommand_from ${{ $parentVarPrefix }}commands" -a {{ .Cmd.Name }} -d "{{ .Cmd.ShortDescription }}"
-
 {{- range .Cmd.Opts }}
-{{- template "opt" (map "Opt" . "Cond" $cond) -}}
+{{- template "opt" (map "Opt" . "ProgName" $progName "VarName" $varName "Cmd" $cmd) -}}
 {{ end -}}
 
+{{ if gt (len .Cmd.Subcommands) 0  }}
+set -l {{ $varPrefix }}commands {{- range .Cmd.Subcommands }} {{ .Name }}{{ end -}}
 {{ $cmdName := .Cmd.Name }}
 {{- range .Cmd.Subcommands }}
-{{- template "cmd" (map "Cmd" . "Prefix" (join $cond "; ") "VarPrefix" $varPrefix) -}}
+{{- template "cmd" (map "Cmd" . "ProgName" $progName "VarName" $varName) -}}
 {{ end -}}
+{{- end -}}
 {{ end }}
 
 {{- /* Top-level commands */ -}}
-{{ range .RootCmd.Subcommands }}
-{{- template "cmd" (map "Cmd" .) -}}
+{{ range $rootCmd.Subcommands }}
+{{- template "cmd" (map "Cmd" . "ProgName" $progName "VarName" $varName) -}}
 {{ end }}
 
 {{- /* Root command options */ -}}
-{{ range .RootCmd.Opts }}
-{{- template "opt" (map "Opt" . "Cond" "not __fish_seen_subcommand_from $commands") -}}
+{{ range $rootCmd.Opts }}
+{{- template "opt" (map "Opt" . "ProgName" $progName "VarName" $varName "Cmd" $rootCmd) -}}
 {{ end }}
 
 {{- /* Global options */ -}}
 {{ range .GlobalOpts }}
-{{- template "opt" (map "Opt" .) -}}
+{{- template "opt" (map "Opt" . "ProgName" $progName "VarName" $varName) -}}
 {{ end }}
 `))
diff --git a/tpl_funcs.go b/tpl_funcs.go
index ae08936..dd41bfa 100644
--- a/tpl_funcs.go
+++ b/tpl_funcs.go
@@ -7,8 +7,10 @@ import (
 )
 
 var tplFuncs = template.FuncMap{
-	"map":  tplMap,
-	"join": tplJoin,
+	"map":       tplMap,
+	"join":      tplJoin,
+	"under":     tplUnder,
+	"varPrefix": tplVarPrefix,
 }
 
 func tplMap(vals ...any) (map[string]any, error) {
@@ -26,3 +28,15 @@ func tplMap(vals ...any) (map[string]any, error) {
 func tplJoin(strs ...string) string {
 	return strings.Join(strs, "")
 }
+
+func tplUnder(s string) string {
+	return strings.ToLower(strings.ReplaceAll(strings.ReplaceAll(s, " ", "_"), "-", "_"))
+}
+
+func tplVarPrefix(s string) string {
+	if s == "" {
+		return ""
+	}
+
+	return tplUnder(s) + "_"
+}