From 9b9e12b4b9b21ed76fcac76594e8739da32035b4 Mon Sep 17 00:00:00 2001
From: Nick Lundin <nick@moontower.net>
Date: Mon, 23 Sep 2024 22:14:14 -0500
Subject: [PATCH] Added some status line info to the submodule list view

---
 pkg/commands/git_commands/submodule.go  | 67 +++++++++++++++++++++++++
 pkg/commands/models/submodule_config.go | 10 ++--
 pkg/gui/presentation/submodules.go      | 32 +++++++++++-
 3 files changed, 105 insertions(+), 4 deletions(-)

diff --git a/pkg/commands/git_commands/submodule.go b/pkg/commands/git_commands/submodule.go
index 40a0d35098f..ac0242db6fb 100644
--- a/pkg/commands/git_commands/submodule.go
+++ b/pkg/commands/git_commands/submodule.go
@@ -54,6 +54,62 @@ func (self *SubmoduleCommands) GetConfigs(parentModule *models.SubmoduleConfig)
 		}
 	}
 
+	// TODO: unsure if this is the best way to go about this.
+	// TODO: might be better to do a single `git submodule foreach git status` and parse the Output
+	getSubmoduleHead := func(cmds *SubmoduleCommands, path string) (string, error) {
+		// git -C </path/to/submodule> symbolic-ref -q HEAD
+		// Output is empty in detached state
+		// Output if on a branch: refs/heads/main
+		cmdArgs := NewGitCmd("symbolic-ref").
+			Dir(path).
+			Arg("-q", "HEAD").
+			ToArgv()
+
+		return cmds.cmd.New(cmdArgs).RunWithOutput()
+	}
+
+	formatSubmoduleHead := func(head string, prefix string) string {
+		// refs/heads/main ---> main
+		if strings.HasPrefix(head, prefix) {
+			return head[len(prefix) : len(head)-1]
+		}
+		return head
+	}
+
+	getPorcelainStatus := func(cmds *SubmoduleCommands, path string) (int, int, int, error) {
+		cmdArgs := NewGitCmd("status").
+			Dir(path).
+			Arg("--porcelain").
+			ToArgv()
+
+		output, err := cmds.cmd.New(cmdArgs).RunWithOutput()
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		lines := strings.Split(output, "\n")
+
+		stagedCount := 0
+		unstagedCount := 0
+		untrackedCount := 0
+
+		for _, line := range lines {
+			if len(line) == 0 {
+				continue
+			}
+			firstChar := line[0:1]
+			switch firstChar {
+			case " ":
+				unstagedCount++
+			case "?":
+				untrackedCount++
+			default:
+				stagedCount++
+			}
+
+		}
+		return stagedCount, unstagedCount, untrackedCount, nil
+	}
+
 	configs := []*models.SubmoduleConfig{}
 	lastConfigIdx := -1
 	for scanner.Scan() {
@@ -70,6 +126,17 @@ func (self *SubmoduleCommands) GetConfigs(parentModule *models.SubmoduleConfig)
 		if lastConfigIdx != -1 {
 			if path, ok := firstMatch(line, `\s*path\s*=\s*(.*)\s*`); ok {
 				configs[lastConfigIdx].Path = path
+				head, err := getSubmoduleHead(self, path)
+				if err == nil {
+					formattedHead := formatSubmoduleHead(head, "refs/heads/")
+					configs[lastConfigIdx].Head = strings.TrimSpace(formattedHead)
+				}
+				stagedCount, unstagedCount, untrackedCount, err := getPorcelainStatus(self, path)
+				if err == nil {
+					configs[lastConfigIdx].NumStagedFiles = stagedCount
+					configs[lastConfigIdx].NumUnstagedChanges = unstagedCount
+					configs[lastConfigIdx].NumUntrackedChanges = untrackedCount
+				}
 				nestedConfigs, err := self.GetConfigs(configs[lastConfigIdx])
 				if err == nil {
 					configs = append(configs, nestedConfigs...)
diff --git a/pkg/commands/models/submodule_config.go b/pkg/commands/models/submodule_config.go
index 7df0d131abf..1919b12cafa 100644
--- a/pkg/commands/models/submodule_config.go
+++ b/pkg/commands/models/submodule_config.go
@@ -3,9 +3,13 @@ package models
 import "path/filepath"
 
 type SubmoduleConfig struct {
-	Name string
-	Path string
-	Url  string
+	Name                string
+	Path                string
+	Url                 string
+	Head                string
+	NumStagedFiles      int
+	NumUnstagedChanges  int
+	NumUntrackedChanges int
 
 	ParentModule *SubmoduleConfig // nil if top-level
 }
diff --git a/pkg/gui/presentation/submodules.go b/pkg/gui/presentation/submodules.go
index 72c6bfc081a..bb9d8a101b6 100644
--- a/pkg/gui/presentation/submodules.go
+++ b/pkg/gui/presentation/submodules.go
@@ -1,6 +1,8 @@
 package presentation
 
 import (
+	"fmt"
+
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/theme"
 	"github.com/samber/lo"
@@ -13,7 +15,11 @@ func GetSubmoduleListDisplayStrings(submodules []*models.SubmoduleConfig) [][]st
 }
 
 func getSubmoduleDisplayStrings(s *models.SubmoduleConfig) []string {
-	name := s.Name
+	// Pad right with some spaces to the end of the HEAD so that it (hopefully) aligns well.
+	// Put the HEAD first because those are more likely to be similar lengths than the repo name.
+	name := fmt.Sprintf("%-20s %s", s.Head,
+		s.Name,
+	)
 	if s.ParentModule != nil {
 		indentation := ""
 		for p := s.ParentModule; p != nil; p = p.ParentModule {
@@ -23,5 +29,29 @@ func getSubmoduleDisplayStrings(s *models.SubmoduleConfig) []string {
 		name = indentation + "- " + s.Name
 	}
 
+	if s.NumStagedFiles != 0 {
+		name = fmt.Sprintf(
+			"%s +%d",
+			name,
+			s.NumStagedFiles,
+		)
+	}
+
+	if s.NumUnstagedChanges != 0 {
+		name = fmt.Sprintf(
+			"%s !%d",
+			name,
+			s.NumUnstagedChanges,
+		)
+	}
+
+	if s.NumUntrackedChanges != 0 {
+		name = fmt.Sprintf(
+			"%s ?%d ",
+			name,
+			s.NumUntrackedChanges,
+		)
+	}
+
 	return []string{theme.DefaultTextColor.Sprint(name)}
 }