Merge pull request #1 from akupila/fix-zsh-width

Add flag to output zsh glitch (%G) characters

README.md

@@ -221,19 +221,21 @@ The code has no vendored dependencies so no need to worry about that.
 Execute `gitprompt` as part of `PROMPT`. Add this to your  `~/.zshrc`:
 
 ```
-export PROMPT='$PROMPT $(gitprompt)'
+export PROMPT='$PROMPT %{$(gitprompt -zsh)%}'
 ```
 
+> The `-zsh` flag makes `gitprompt` output the correct width of visible
+> characters, which fixes counting ansi color codes (breaks wrapping).
+
 Now reload the config (`source ~/.zshrc`) and gitprompt should show up. Feel
 free to add anything else here too, just execute `gitprompt` where you want the
 status, for example _(this was used for taking the screenshots in the readme)_:
 
 ```
-local ret_status="%(?:%{$fg_bold[green]%}›:%{$fg_bold[red]%}›)"
-local dir="%{$fg[cyan]%}%3d%{$reset_color%}"
-export PROMPT='${ret_status} ${dir} $(gitprompt)'
+export PROMPT='%(?:%{$fg_bold[green]%}›:%{$fg_bold[red]%}›) %{$fg[cyan]%}%3d %{$(gitprompt -zsh)%}%{$reset_color%}'
 ```
 
+
 Alternatively, you can add this to `RPROMPT` instead, which will make the
 status appear on the right hand side of the screen. `gitprompt` will by default
 add a trailing space so you you may want to customize the formatting if you

cmd/gitprompt/gitprompt.go

@@ -4,7 +4,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"strings"
 
 	"github.com/akupila/gitprompt"
 )
@@ -55,8 +54,8 @@ var exampleStatus = &gitprompt.GitStatus{
 }
 
 var formatHelp = func() string {
-	return strings.TrimSpace(fmt.Sprintf(`
-Define output format.
+	example, _ := gitprompt.Print(exampleStatus, defaultFormat)
+	return fmt.Sprintf(`Define output format.
 
 Default format is: %q
 Example result:    %s
@@ -94,12 +93,12 @@ Text attributes:
 	@f	Set faint/dim color
 	@F	Clear faint/dim color
 	@i	Set italic
-	@I	Clear italic
-`, defaultFormat, gitprompt.Print(exampleStatus, defaultFormat)))
+	@I	Clear italic`, defaultFormat, example)
 }
 
 func main() {
 	v := flag.Bool("version", false, "Print version inforformation.")
+	zsh := flag.Bool("zsh", false, "Print zsh width control characters")
 	flag.Var(&format, "format", formatHelp())
 	flag.Parse()
 
@@ -111,5 +110,14 @@ func main() {
 		os.Exit(0)
 	}
 
-	gitprompt.Exec(format.String())
+	s, err := gitprompt.Parse()
+	if err != nil {
+		_, _ = fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	out, num := gitprompt.Print(s, format.String())
+	_, _ = fmt.Fprint(os.Stdout, out)
+	if *zsh {
+		_, _ = fmt.Fprintf(os.Stdout, "%%%dG", num)
+	}
 }

parser.go

@@ -9,6 +9,18 @@ import (
 	"strings"
 )
 
+// GitStatus is the parsed status for the current state in git.
+type GitStatus struct {
+	Sha       string
+	Branch    string
+	Untracked int
+	Modified  int
+	Staged    int
+	Conflicts int
+	Ahead     int
+	Behind    int
+}
+
 // Parse parses the status for the repository from git. Returns nil if the
 // current directory is not part of a git repository.
 func Parse() (*GitStatus, error) {

printer.go

@@ -69,12 +69,15 @@ type group struct {
 
 	hasData  bool
 	hasValue bool
+	width    int
 }
 
 // Print prints the status according to the format.
-func Print(s *GitStatus, format string) string {
+//
+// The integer returned is the print width of the string.
+func Print(s *GitStatus, format string) (string, int) {
 	if s == nil {
-		return ""
+		return "", 0
 	}
 
 	in := make(chan rune)
@@ -93,7 +96,7 @@ func Print(s *GitStatus, format string) string {
 	return buildOutput(s, in)
 }
 
-func buildOutput(s *GitStatus, in chan rune) string {
+func buildOutput(s *GitStatus, in chan rune) (string, int) {
 	root := &group{}
 	g := root
 
@@ -148,6 +151,7 @@ func buildOutput(s *GitStatus, in chan rune) string {
 				g.parent.format = g.format
 				g.parent.format.setColor(0)
 				g.parent.format.clearAttributes()
+				g.parent.width += g.width
 			}
 			g = g.parent
 		default:
@@ -170,7 +174,7 @@ func buildOutput(s *GitStatus, in chan rune) string {
 	g.format.clearAttributes()
 	g.format.printANSI(&g.buf)
 
-	return root.buf.String()
+	return root.buf.String(), root.width
 }
 
 func setColor(g *group, ch rune) {
@@ -274,15 +278,16 @@ func (g *group) addRune(r rune) {
 	if !unicode.IsSpace(r) {
 		g.format.printANSI(&g.buf)
 	}
+	g.width++
 	g.buf.WriteRune(r)
 }
 
 func (g *group) addString(s string) {
 	g.format.printANSI(&g.buf)
+	g.width += len(s)
 	g.buf.WriteString(s)
 }
 
 func (g *group) addInt(i int) {
-	g.format.printANSI(&g.buf)
-	g.buf.WriteString(strconv.Itoa(i))
+	g.addString(strconv.Itoa(i))
 }