pest.go 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. package main
  2. import (
  3. "fmt"
  4. "github.com/go-git/go-git/v5"
  5. "github.com/gookit/color"
  6. "github.com/pelletier/go-toml/v2"
  7. "os"
  8. "os/user"
  9. "strings"
  10. )
  11. // go-pest -- A reimplementation of cl-pest in Go, for the sake of learning a bit of Go
  12. // See: https://git.spwbk.site/swatson/cl-pest/src/master/pest.lisp
  13. type PestConfig struct {
  14. Git GitConfig `toml:"git"`
  15. Prompt PromptConfig `toml:"prompt"`
  16. }
  17. type GitConfig struct {
  18. DisplayHead bool `toml:"display_head"`
  19. DisplayBranch bool `toml:"display_branch"`
  20. GitPrefix string `toml:"git_prefix"`
  21. Colors Colors `toml:"colors"`
  22. }
  23. type PromptConfig struct {
  24. DisplayUser bool `toml:"display_user"`
  25. UserSuffix string `toml:"user_suffix"`
  26. DisplayHostname bool `toml:"display_hostname"`
  27. HostnameSuffix string `toml:"hostname_suffix"`
  28. DisplayPwd bool `toml:"display_pwd"`
  29. PwdSuffix string `toml:"pwd_suffix"`
  30. PromptChar string `toml:"prompt_char"`
  31. Colors Colors `toml:"colors"`
  32. }
  33. type Colors struct {
  34. Fg []int `toml:"fg"`
  35. Bg []int `toml:"bg"`
  36. }
  37. const tomlData = `
  38. [git]
  39. display_head = false
  40. display_branch = false
  41. git_prefix = ""
  42. [git.colors]
  43. fg = [0, 120, 50]
  44. bg = [0, 0, 0]
  45. [prompt]
  46. display_user = false
  47. user_suffix = ""
  48. display_hostname = false
  49. hostname_suffix = ""
  50. display_pwd = true
  51. pwd_suffix = ""
  52. prompt_char = " λ "
  53. [prompt.colors]
  54. fg = [255, 255, 255]
  55. bg = [0, 0, 0]
  56. `
  57. func parseTomlFromString(tomlString string) (PestConfig, error) {
  58. var config PestConfig
  59. err := toml.Unmarshal([]byte(tomlString), &config)
  60. return config, err
  61. }
  62. func parseTomlFromFile(filePath string) (PestConfig, error) {
  63. var cfg PestConfig
  64. fileContent, err := os.ReadFile(filePath)
  65. switch err {
  66. case nil:
  67. {
  68. err = toml.Unmarshal(fileContent, &cfg)
  69. return cfg, err
  70. }
  71. default:
  72. {
  73. // We had some kind of error, use the builtin config
  74. cfg, err = parseTomlFromString(tomlData)
  75. return cfg, err
  76. }
  77. }
  78. }
  79. func gitGetHead(filePath string) (string, string) {
  80. gitObj, err := git.PlainOpen(filePath)
  81. if err != nil {
  82. // Likely not a git dir, return empty
  83. return "", ""
  84. }
  85. headRef, err := gitObj.Head()
  86. if err != nil {
  87. panic(err)
  88. }
  89. sha := headRef.Hash().String()[:6]
  90. branch := headRef.Name().Short()
  91. return sha, branch
  92. }
  93. func getHostname() string {
  94. hostname, err := os.Hostname()
  95. if err != nil {
  96. panic(err)
  97. }
  98. return hostname
  99. }
  100. func getUser() string {
  101. username, err := user.Current()
  102. if err != nil {
  103. panic(err)
  104. }
  105. return username.Username
  106. }
  107. func getPwd() string {
  108. pwd, err := os.Getwd()
  109. if err != nil {
  110. panic(err)
  111. }
  112. return pwd
  113. }
  114. type RGBColorInput struct {
  115. r uint8
  116. g uint8
  117. b uint8
  118. }
  119. // make_style -- given a set of RGBColorInputs return a color.RGBColor pair
  120. func make_style(fg, bg RGBColorInput) color.Style {
  121. fg_s := color.RGB(fg.r, fg.g, fg.b)
  122. bg_s := color.RGB(bg.r, bg.g, bg.b, true)
  123. style := color.New(fg_s.Color(), bg_s.Color())
  124. return style
  125. }
  126. func assemble_prompt(config *PestConfig) {
  127. var prompt_builder strings.Builder
  128. var git_prompt_builder strings.Builder
  129. if config.Git.DisplayHead || config.Git.DisplayBranch {
  130. sha, branch := gitGetHead(getPwd())
  131. if config.Git.DisplayHead && config.Git.DisplayBranch {
  132. git_prompt := sha + "|" + branch
  133. git_prompt_builder.WriteString(config.Git.GitPrefix)
  134. git_prompt_builder.WriteString(git_prompt)
  135. } else if config.Git.DisplayHead {
  136. git_prompt_builder.WriteString(config.Git.GitPrefix)
  137. git_prompt_builder.WriteString(sha)
  138. } else if config.Git.DisplayBranch {
  139. git_prompt_builder.WriteString(config.Git.GitPrefix)
  140. git_prompt_builder.WriteString(branch)
  141. }
  142. git_prompt_builder.WriteString(" ")
  143. }
  144. git_prompt_color := make_style(
  145. RGBColorInput{
  146. uint8(config.Git.Colors.Fg[0]),
  147. uint8(config.Git.Colors.Fg[1]),
  148. uint8(config.Git.Colors.Fg[2]),
  149. },
  150. RGBColorInput{
  151. uint8(config.Git.Colors.Bg[0]),
  152. uint8(config.Git.Colors.Bg[1]),
  153. uint8(config.Git.Colors.Bg[2]),
  154. })
  155. if config.Prompt.DisplayUser {
  156. prompt_builder.WriteString(getUser())
  157. prompt_builder.WriteString(config.Prompt.UserSuffix)
  158. }
  159. if config.Prompt.DisplayHostname {
  160. prompt_builder.WriteString(getHostname())
  161. prompt_builder.WriteString(config.Prompt.HostnameSuffix)
  162. }
  163. if config.Prompt.DisplayPwd {
  164. prompt_builder.WriteString(getPwd())
  165. prompt_builder.WriteString(config.Prompt.PwdSuffix)
  166. }
  167. prompt_builder.WriteString(config.Prompt.PromptChar)
  168. prompt_color := make_style(
  169. RGBColorInput{
  170. uint8(config.Prompt.Colors.Fg[0]),
  171. uint8(config.Prompt.Colors.Fg[1]),
  172. uint8(config.Prompt.Colors.Fg[2]), // This will of course panic if this number is > 255
  173. },
  174. RGBColorInput{
  175. uint8(config.Prompt.Colors.Bg[0]),
  176. uint8(config.Prompt.Colors.Bg[1]),
  177. uint8(config.Prompt.Colors.Bg[2]),
  178. })
  179. git_prompt_color.Printf(git_prompt_builder.String())
  180. prompt_color.Printf(prompt_builder.String())
  181. }
  182. func main() {
  183. toml, err := parseTomlFromFile("/home/swatson/Repos/go-pest/pest.cfg")
  184. if err != nil {
  185. fmt.Println("Couldn't parse TOML const!")
  186. }
  187. assemble_prompt(&toml)
  188. }