more complicated compgen for directory vs file completion

pkg/server/server.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"sort"
 	"strings"
 	"sync"
 
@@ -51,18 +52,15 @@ func (m *MServer) ProcessCommandPacket(pk packet.CommandPacketType) {
 	return
 }
 
-func (m *MServer) runCompGen(compPk *packet.CompGenPacketType) {
+func runSingleCompGen(cwd string, compType string, prefix string) ([]string, bool, error) {
-	reqId := compPk.GetReqId()
+	if !packet.IsValidCompGenType(compType) {
-	if !packet.IsValidCompGenType(compPk.CompType) {
+		return nil, false, fmt.Errorf("invalid compgen type '%s'", compType)
-		m.Sender.SendErrorResponse(reqId, fmt.Errorf("invalid compgen type '%s'", compPk.CompType))
-		return
 	}
-	compGenCmdStr := fmt.Sprintf("cd %s; compgen -A %s -- %s | sort | uniq | head -n %d", shellescape.Quote(compPk.Cwd), shellescape.Quote(compPk.CompType), shellescape.Quote(compPk.Prefix), packet.MaxCompGenValues+1)
+	compGenCmdStr := fmt.Sprintf("cd %s; compgen -A %s -- %s | sort | uniq | head -n %d", shellescape.Quote(cwd), shellescape.Quote(compType), shellescape.Quote(prefix), packet.MaxCompGenValues+1)
 	ecmd := exec.Command("bash", "-c", compGenCmdStr)
 	outputBytes, err := ecmd.Output()
 	if err != nil {
-		m.Sender.SendErrorResponse(reqId, fmt.Errorf("compgen error: %w", err))
+		return nil, false, fmt.Errorf("compgen error: %w", err)
-		return
 	}
 	outputStr := string(outputBytes)
 	parts := strings.Split(outputStr, "\n")
@@ -74,7 +72,68 @@ func (m *MServer) runCompGen(compPk *packet.CompGenPacketType) {
 		hasMore = true
 		parts = parts[0:packet.MaxCompGenValues]
 	}
-	m.Sender.SendResponse(reqId, map[string]interface{}{"comps": parts, "hasmore": hasMore})
+	return parts, hasMore, nil
+}
+
+func appendSlashes(comps []string) {
+	for idx, comp := range comps {
+		comps[idx] = comp + "/"
+	}
+}
+
+func strArrToMap(strs []string) map[string]bool {
+	rtn := make(map[string]bool)
+	for _, s := range strs {
+		rtn[s] = true
+	}
+	return rtn
+}
+
+func (m *MServer) runFileCompGen(compPk *packet.CompGenPacketType) {
+	// get directories and files, unique them and put slashes on directories for completion
+	reqId := compPk.GetReqId()
+	compDirs, hasMoreDirs, err := runSingleCompGen(compPk.Cwd, "directory", compPk.Prefix)
+	if err != nil {
+		m.Sender.SendErrorResponse(reqId, err)
+		return
+	}
+	compFiles, hasMoreFiles, err := runSingleCompGen(compPk.Cwd, "file", compPk.Prefix)
+	if err != nil {
+		m.Sender.SendErrorResponse(reqId, err)
+		return
+	}
+
+	dirMap := strArrToMap(compDirs)
+	// seed comps with dirs (but append slashes)
+	comps := compDirs
+	appendSlashes(comps)
+	// add files that are not directories (look up in dirMap)
+	for _, file := range compFiles {
+		if dirMap[file] {
+			continue
+		}
+		comps = append(comps, file)
+	}
+	sort.Strings(comps) // resort
+	m.Sender.SendResponse(reqId, map[string]interface{}{"comps": comps, "hasmore": (hasMoreFiles || hasMoreDirs)})
+	return
+}
+
+func (m *MServer) runCompGen(compPk *packet.CompGenPacketType) {
+	reqId := compPk.GetReqId()
+	if compPk.CompType == "file" {
+		m.runFileCompGen(compPk)
+		return
+	}
+	comps, hasMore, err := runSingleCompGen(compPk.Cwd, compPk.CompType, compPk.Prefix)
+	if err != nil {
+		m.Sender.SendErrorResponse(reqId, err)
+		return
+	}
+	if compPk.CompType == "directory" {
+		appendSlashes(comps)
+	}
+	m.Sender.SendResponse(reqId, map[string]interface{}{"comps": comps, "hasmore": hasMore})
 	return
 }