pages/blog/go-shell-prompt.md (view raw)
1---
2template:
3slug: go-shell-prompt
4title: Writing a shell prompt in Go
5subtitle: Kinda faster than bash
6date: 2021-08-12
7---
8
9For context, my bash prompt was previously [written in, well,
10bash](https://git.icyphox.sh/dotfiles/tree/bash/.bashrc.d/99-prompt.bash?id=d7b391845abc7e97f2b1b96c34b4b1789b2ab541).
11It used to call out to `git` for getting the branch and worktree status
12info. Parsing the output of `git status` and all that. It was ok, but I
13wanted something ... cleaner.
14
15I chose Go, despite having written
16[nicy](https://github.com/icyphox/nicy) in Nim[^1]; I'm in a Go-phase right
17now, just like I was in a Nim-phase back in 2018. Anyway, let's cut to
18the chase.
19
20[^1]: It's a prompt "framework" thing that I actually only used for a
21 month or so.
22
23## the basics
24
25The current working directory is the bare minimum in a prompt. I prefer
26having it shortened; for example: `/home/icy/docs/books/foo.epub` →
27`~/d/b/foo.epub`. Let's write a function `trimPath` to do this for us:
28
29```go
30// Truncates the current working directory:
31// /home/icy/foo/bar -> ~/f/bar
32func trimPath(cwd, home string) string {
33 var path string
34 if strings.HasPrefix(cwd, home) {
35 path = "~" + strings.TrimPrefix(cwd, home)
36 } else {
37 // If path doesn't contain $HOME, return the
38 // entire path as is.
39 path = cwd
40 return path
41 }
42 items := strings.Split(path, "/")
43 truncItems := []string{}
44 for i, item := range items {
45 if i == (len(items) - 1) {
46 truncItems = append(truncItems, item)
47 break
48 }
49 truncItems = append(truncItems, item[:1])
50 }
51 return filepath.Join(truncItems...)
52}
53```
54
55`trimPath` takes two args: the current working directory `cwd`, and the
56home directory `home`. We first check if `cwd` starts with `home`, i.e.
57we're in a subdirectory of `home`; if yes, trim `home` from `cwd`, and
58replace it with a tilde `~`. We now have `~/docs/books/foo.epub`.
59
60Also note that we return the path as-is if we're not in a subdir of
61`home` -- i.e. paths under `/`, like `/usr`, etc. I like to see these
62completely, just to be sure.
63
64We then split the path at `/`[^2], and truncate each item in the
65resulting list -- except for the last -- down to the first character.
66Join it all together and return the resulting string -- we have
67`~/d/b/foo.epub`.
68
69[^2]: I don't care about Windows.
70
71Next up: color.
72
73```go
74var (
75 red = color("\033[31m%s\033[0m")
76 green = color("\033[32m%s\033[0m")
77 cyan = color("\033[36m%s\033[0m")
78)
79
80func color(s string) func(...interface{}) string {
81 return func(args ...interface{}) string {
82 return fmt.Sprintf(s, fmt.Sprint(args...))
83 }
84}
85```
86... I'll just let you figure this one out.
87
88## git branch and clean/dirty info
89
90The defacto lib for git in Go is often
91[go-git](https://github.com/go-git/go-git). I don't disagree that it is
92a good library: clean APIs, good docs. It just has one huge issue, and
93especially so in our case. It's `worktree.Status()` function -- used to
94fetch the worktree status -- is [awfully
95slow](https://github.com/go-git/go-git/issues/327). It's not noticeable
96in small repositories, but even relatively large ones (~30MB) tend to
97take about 20 seconds. That's super not ideal for a prompt.[^3]
98
99[^3]: https://git.icyphox.sh/dotfiles/commit/?id=e1f6aaaf6ffd35224b5d3f057c28fb2560e1c3b0
100
101The alternative? [libgit2/git2go](https://github.com/libgit2/git2go), of
102course! It's the Go bindings for libgit2 -- obviously, it requires CGo,
103but who cares.
104
105First things first, let's write `getGitDir` to find `.git` (indicating
106that it's a git repo), and return the repository path. We'll need this
107to use git2go's `OpenRepository()`.
108
109```go
110// Recursively traverse up until we find .git
111// and return the git repo path.
112func getGitDir() string {
113 cwd, _ := os.Getwd()
114 for {
115 dirs, _ := os.ReadDir(cwd)
116 for _, d := range dirs {
117 if ".git" == d.Name() {
118 return cwd
119 } else if cwd == "/" {
120 return ""
121 }
122 }
123 cwd = filepath.Dir(cwd)
124 }
125}
126```
127
128This traverses up parent directories until it finds `.git`, else,
129returns an empty string if we've reached `/`. For example: if you're in
130`~/code/foo/bar`, and the git repo root is at `~/code/foo/.git`, this
131function will find it.
132
133Alright, let's quickly write two more functions to return the git branch
134name, and the repository status -- i.e., dirty or clean.
135
136```go
137// Returns the current git branch or current ref sha.
138func gitBranch(repo *git.Repository) string {
139 ref, _ := repo.Head()
140 if ref.IsBranch() {
141 name, _ := ref.Branch().Name()
142 return name
143 } else {
144 return ref.Target().String()[:7]
145 }
146}
147```
148
149This takes a `*git.Repository`, where `git` is `git2go`. We first get
150the `git.Reference` and check whether it's a branch. If yes, return the
151name of the branch, else -- like in the case of a detached HEAD state --
152we just return a short hash.
153
154```go
155// Returns • if clean, else ×.
156func gitStatus(repo *git.Repository) string {
157 sl, _ := repo.StatusList(&git.StatusOptions{
158 Show: git.StatusShowIndexAndWorkdir,
159 Flags: git.StatusOptIncludeUntracked,
160 })
161 n, _ := sl.EntryCount()
162 if n != 0 {
163 return red("×")
164 } else {
165 return green("•")
166 }
167}
168```
169
170We use the
171[`StatusList`](https://godocs.io/github.com/libgit2/git2go/v31#Repository.StatusList)
172function to produce a `StatusList` object. We then check the
173`EntryCount`, i.e., the number of modified/untracked/etc. files
174contained in `StatusList`.[^4] If this number is 0, our repo is clean;
175dirty otherwise. Colored symbols are returned accordingly.
176
177[^4]: This took me a lot of going back and forth, and reading
178 https://libgit2.org/libgit2/ex/HEAD/status.html to figure out.
179
180## putting it all together
181
182Home stretch. Let's write `makePrompt` to make our prompt.
183
184```go
185const (
186 promptSym = "▲"
187)
188
189func makePrompt() string {
190 cwd, _ := os.Getwd()
191 home := os.Getenv("HOME")
192 gitDir := getGitDir()
193 if len(gitDir) > 0 {
194 repo, _ := git.OpenRepository(getGitDir())
195 return fmt.Sprintf(
196 "\n%s (%s %s)\n%s",
197 cyan(trimPath(cwd, home)),
198 gitBranch(repo),
199 gitStatus(repo),
200 promptSym,
201 )
202 }
203 return fmt.Sprintf(
204 "\n%s\n%s",
205 cyan(trimPath(cwd, home)),
206 promptSym,
207 )
208}
209
210func main() {
211 fmt.Println(makePrompt())
212}
213```
214
215There isn't much going on here. Get the necessary pieces like the
216current working directory, home and the git repo path. We return the
217formatted prompt string according to whether we're in git or not.
218
219Setting the prompt is simple. Point `PS1` to the built binary:
220
221```bash
222PS1='$(~/dotfiles/prompt/prompt) '
223```
224
225And here's what it looks like, rendered:
226
227
228## benchmarking
229
230Both "benchmarks" were run inside a sufficiently large git repository,
231deep inside many subdirs.
232
233To time the old bash prompt, I just copied all the bash functions,
234pasted it in my shell and ran:
235
236```shell
237~/C/d/a/n/y/b/d/t/yaml-1.1 (master •)
238▲ time echo -e "\n$(prompt_pwd)$(git_branch)\n▲$(rootornot)"
239
240# output
241~/C/d/a/n/y/b/d/t/yaml-1.1 (master •)
242▲
243
244real 0m0.125s
245user 0m0.046s
246sys 0m0.079s
247
248```
249
2500.125s. Not too bad. Let's see how long our Go prompt takes.
251
252```shell
253~/C/d/a/n/y/b/d/t/yaml-1.1 (master •)
254▲ time ~/dotfiles/prompt/prompt
255
256# output
257~/C/d/a/n/y/b/d/t/yaml-1.1 (master •)
258▲
259
260real 0m0.074s
261user 0m0.031s
262sys 0m0.041s
263
264```
265
2660.074s! That's pretty fast. I ran these tests a few more times, and the
267bash version was consistently slower -- averaging ~0.120s; the Go
268version averaging ~0.70s. That's a win.
269
270You can find the entire source [here](https://git.icyphox.sh/dotfiles/tree/prompt).