go-pest/pest.go

package main

import (
	"fmt"
	"github.com/go-git/go-git/v5"
	"github.com/gookit/color"
	"github.com/pelletier/go-toml/v2"
	"os"
	"os/user"
	"strings"
)

// go-pest -- A reimplementation of cl-pest in Go, for the sake of learning a bit of Go
// See: https://git.spwbk.site/swatson/cl-pest/src/master/pest.lisp

type PestConfig struct {
	Git    GitConfig    `toml:"git"`
	Prompt PromptConfig `toml:"prompt"`
}

type GitConfig struct {
	DisplayHead   bool   `toml:"display_head"`
	DisplayBranch bool   `toml:"display_branch"`
	GitPrefix     string `toml:"git_prefix"`
	Colors        Colors `toml:"colors"`
}

type PromptConfig struct {
	DisplayUser     bool   `toml:"display_user"`
	UserSuffix      string `toml:"user_suffix"`
	DisplayHostname bool   `toml:"display_hostname"`
	HostnameSuffix  string `toml:"hostname_suffix"`
	DisplayPwd      bool   `toml:"display_pwd"`
	PwdSuffix       string `toml:"pwd_suffix"`
	PromptChar      string `toml:"prompt_char"`
	Colors          Colors `toml:"colors"`
}

type Colors struct {
	Fg []int `toml:"fg"`
	Bg []int `toml:"bg"`
}

const tomlData = `
[git]
display_head = false
display_branch = false
git_prefix = ""
[git.colors]
fg = [0, 120, 50]
bg = [0, 0, 0]

[prompt]
display_user = false
user_suffix = ""
display_hostname = false
hostname_suffix = ""
display_pwd = true
pwd_suffix = ""
prompt_char = " λ "
[prompt.colors]
fg = [255, 255, 255]
bg = [0, 0, 0]
`

func parseTomlFromString(tomlString string) (PestConfig, error) {
	var config PestConfig
	err := toml.Unmarshal([]byte(tomlString), &config)
	return config, err
}

func parseTomlFromFile(filePath string) (PestConfig, error) {
	var cfg PestConfig
	fileContent, err := os.ReadFile(filePath)
	switch err {
	case nil:
		{
			err = toml.Unmarshal(fileContent, &cfg)
			return cfg, err
		}
	default:
		{
			// We had some kind of error, use the builtin config
			cfg, err = parseTomlFromString(tomlData)
			return cfg, err
		}
	}
}

func gitGetHead(filePath string) (string, string) {
	gitObj, err := git.PlainOpen(filePath)
	if err != nil {
		panic(err)
	}

	headRef, err := gitObj.Head()
	if err != nil {
		panic(err)
	}

	sha := headRef.Hash().String()[:6]
	branch := headRef.Name().Short()
	return sha, branch
}

func getHostname() string {
	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}
	return hostname
}

func getUser() string {
	username, err := user.Current()
	if err != nil {
		panic(err)
	}
	return username.Username
}

func getPwd() string {
	pwd, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	return pwd
}

type RGBColorInput struct {
	r uint8
	g uint8
	b uint8
}

// make_style -- given a set of RGBColorInputs return a color.RGBColor pair
func make_style(fg, bg RGBColorInput) color.Style {
	fg_s := color.RGB(fg.r, fg.g, fg.b)
	bg_s := color.RGB(bg.r, bg.g, bg.b, true)
	style := color.New(fg_s.Color(), bg_s.Color())
	return style
}

func assemble_prompt(config *PestConfig) {
	var prompt_builder strings.Builder
	var git_prompt_builder strings.Builder

	if config.Git.DisplayHead || config.Git.DisplayBranch {
		sha, branch := gitGetHead(getPwd())
		if config.Git.DisplayHead && config.Git.DisplayBranch {
			git_prompt := sha + "|" + branch
			git_prompt_builder.WriteString(config.Git.GitPrefix)
			git_prompt_builder.WriteString(git_prompt)
		} else if config.Git.DisplayHead {
			git_prompt_builder.WriteString(config.Git.GitPrefix)
			git_prompt_builder.WriteString(sha)
		} else if config.Git.DisplayBranch {
			git_prompt_builder.WriteString(config.Git.GitPrefix)
			git_prompt_builder.WriteString(branch)
		}
		git_prompt_builder.WriteString(" ")
	}

	git_prompt_color := make_style(
		RGBColorInput{
			uint8(config.Git.Colors.Fg[0]),
			uint8(config.Git.Colors.Fg[1]),
			uint8(config.Git.Colors.Fg[2]),
		},
		RGBColorInput{
			uint8(config.Git.Colors.Bg[0]),
			uint8(config.Git.Colors.Bg[1]),
			uint8(config.Git.Colors.Bg[2]),
		})

	if config.Prompt.DisplayUser {
		prompt_builder.WriteString(getUser())
		prompt_builder.WriteString(config.Prompt.UserSuffix)
	}

	if config.Prompt.DisplayHostname {
		prompt_builder.WriteString(getHostname())
		prompt_builder.WriteString(config.Prompt.HostnameSuffix)
	}

	if config.Prompt.DisplayPwd {
		prompt_builder.WriteString(getPwd())
		prompt_builder.WriteString(config.Prompt.PwdSuffix)
	}

	prompt_builder.WriteString(config.Prompt.PromptChar)

	prompt_color := make_style(
		RGBColorInput{
			uint8(config.Prompt.Colors.Fg[0]),
			uint8(config.Prompt.Colors.Fg[1]),
			uint8(config.Prompt.Colors.Fg[2]), // This will of course panic if this number is > 255
		},
		RGBColorInput{
			uint8(config.Prompt.Colors.Bg[0]),
			uint8(config.Prompt.Colors.Bg[1]),
			uint8(config.Prompt.Colors.Bg[2]),
		})

	git_prompt_color.Printf(git_prompt_builder.String())
	prompt_color.Printf(prompt_builder.String())

}

func main() {

	toml, err := parseTomlFromFile("./pest.cfg")
	if err != nil {
		fmt.Println("Couldn't parse TOML const!")
	}

	assemble_prompt(&toml)

}