refactor: count length using written text

src/color/ansi.go

@@ -7,14 +7,9 @@ import (
 )
 
 const (
-	ansiRegex = "[\u001B\u009B][[\\]()#;?]*(?:(?:(?:[a-zA-Z\\d]*(?:;[a-zA-Z\\d]*)*)?\u0007)|(?:(?:\\d{1,4}(?:;\\d{0,4})*)?[\\dA-PRZcf-ntqry=><~]))"
-
 	zsh  = "zsh"
 	bash = "bash"
 	pwsh = "pwsh"
-
-	str = "STR"
-	url = "URL"
 )
 
 type Ansi struct {
@@ -127,27 +122,6 @@ func (a *Ansi) Init(shell string) {
 	a.shellReservedKeywords = append(a.shellReservedKeywords, shellKeyWordReplacement{"`", "'"})
 }
 
-func (a *Ansi) LenWithoutANSI(text string) int {
-	if len(text) == 0 {
-		return 0
-	}
-	// replace hyperlinks(file/http/https)
-	matches := regex.FindAllNamedRegexMatch(`(?P<STR>\x1b]8;;(file|http|https):\/\/(.+?)\x1b\\(?P<URL>.+?)\x1b]8;;\x1b\\)`, text)
-	for _, match := range matches {
-		text = strings.ReplaceAll(text, match[str], match[url])
-	}
-	// replace console title
-	matches = regex.FindAllNamedRegexMatch(`(?P<STR>\x1b\]0;(.+)\007)`, text)
-	for _, match := range matches {
-		text = strings.ReplaceAll(text, match[str], "")
-	}
-	stripped := regex.ReplaceAllString(ansiRegex, text, "")
-	stripped = strings.ReplaceAll(stripped, a.escapeLeft, "")
-	stripped = strings.ReplaceAll(stripped, a.escapeRight, "")
-	runeText := []rune(stripped)
-	return len(runeText)
-}
-
 func (a *Ansi) generateHyperlink(text string) string {
 	// hyperlink matching
 	results := regex.FindNamedRegexMatch("(?P<all>(?:\\[(?P<name>.+)\\])(?:\\((?P<url>.*)\\)))", text)
@@ -183,8 +157,8 @@ func (a *Ansi) CarriageForward() string {
 	return fmt.Sprintf(a.right, 1000)
 }
 
-func (a *Ansi) GetCursorForRightWrite(text string, offset int) string {
+func (a *Ansi) GetCursorForRightWrite(length, offset int) string {
-	strippedLen := a.LenWithoutANSI(text) + -offset
+	strippedLen := length + (-offset)
 	return fmt.Sprintf(a.left, strippedLen)
 }
 

src/color/ansi_test.go

@@ -6,24 +6,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestLenWithoutAnsi(t *testing.T) {
-	cases := []struct {
-		Text      string
-		ShellName string
-		Expected  int
-	}{
-		{Text: "%{\x1b[44m%}hello%{\x1b[0m%}", ShellName: zsh, Expected: 5},
-		{Text: "\x1b[44mhello\x1b[0m", ShellName: pwsh, Expected: 5},
-		{Text: "\\[\x1b[44m\\]hello\\[\x1b[0m\\]", ShellName: bash, Expected: 5},
-	}
-	for _, tc := range cases {
-		a := Ansi{}
-		a.Init(tc.ShellName)
-		strippedLength := a.LenWithoutANSI(tc.Text)
-		assert.Equal(t, 5, strippedLength)
-	}
-}
-
 func TestGenerateHyperlinkNoUrl(t *testing.T) {
 	cases := []struct {
 		Text      string

src/color/plain_writer.go

@@ -8,6 +8,7 @@ import (
 // PlainWriter writes a plain string
 type PlainWriter struct {
 	builder strings.Builder
+	length  int
 }
 
 func (a *PlainWriter) SetColors(background, foreground string)       {}
@@ -19,6 +20,7 @@ func (a *PlainWriter) Write(background, foreground, text string) {
 		return
 	}
 	writeAndRemoveText := func(text, textToRemove, parentText string) string {
+		a.length += measureText(text)
 		a.builder.WriteString(text)
 		return strings.Replace(parentText, textToRemove, "", 1)
 	}
@@ -30,11 +32,12 @@ func (a *PlainWriter) Write(background, foreground, text string) {
 		text = writeAndRemoveText(textBeforeColorOverride, textBeforeColorOverride, text)
 		text = writeAndRemoveText(innerText, escapedTextSegment, text)
 	}
+	a.length += measureText(text)
 	a.builder.WriteString(text)
 }
 
-func (a *PlainWriter) String() string {
+func (a *PlainWriter) String() (string, int) {
-	return a.builder.String()
+	return a.builder.String(), a.length
 }
 
 func (a *PlainWriter) Reset() {

src/color/text.go

@@ -0,0 +1,8 @@
+package color
+
+import "unicode/utf8"
+
+func measureText(text string) int {
+	length := utf8.RuneCountInString(text)
+	return length
+}

src/color/writer.go

@@ -12,7 +12,7 @@ const (
 
 type Writer interface {
 	Write(background, foreground, text string)
-	String() string
+	String() (string, int)
 	Reset()
 	SetColors(background, foreground string)
 	SetParentColors(background, foreground string)
@@ -28,6 +28,7 @@ type AnsiWriter struct {
 	AnsiColors         AnsiColors
 
 	builder strings.Builder
+	length  int
 }
 
 type Color struct {
@@ -104,6 +105,7 @@ func (a *AnsiWriter) writeColoredText(background, foreground AnsiColor, text str
 	if text == "" || (foreground.IsTransparent() && background.IsTransparent()) {
 		return
 	}
+	a.length += measureText(text)
 	// default to white fg if empty, empty backgrond is supported
 	if foreground.IsEmpty() {
 		foreground = a.getAnsiFromColorString("white", false)
@@ -226,10 +228,11 @@ func (a *AnsiWriter) expandKeyword(keyword string) string {
 	return keyword
 }
 
-func (a *AnsiWriter) String() string {
+func (a *AnsiWriter) String() (string, int) {
-	return a.builder.String()
+	return a.builder.String(), a.length
 }
 
 func (a *AnsiWriter) Reset() {
+	a.length = 0
 	a.builder.Reset()
 }

src/color/writer_test.go

@@ -181,7 +181,7 @@ func TestWriteANSIColors(t *testing.T) {
 			AnsiColors:         &DefaultColors{},
 		}
 		renderer.Write(tc.Colors.Background, tc.Colors.Foreground, tc.Input)
-		got := renderer.String()
+		got, _ := renderer.String()
 		assert.Equal(t, tc.Expected, got, tc.Case)
 	}
 }

src/engine/block.go

@@ -42,6 +42,7 @@ type Block struct {
 	previousActiveSegment *Segment
 	activeBackground      string
 	activeForeground      string
+	length                int
 }
 
 func (b *Block) init(env environment.Environment, writer color.Writer, ansi *color.Ansi) {
@@ -92,7 +93,7 @@ func (b *Block) renderSegmentsText() {
 	}
 }
 
-func (b *Block) renderSegments() string {
+func (b *Block) renderSegments() (string, int) {
 	defer b.writer.Reset()
 	for _, segment := range b.Segments {
 		if !segment.active {
@@ -180,7 +181,7 @@ func (b *Block) debug() (int, []*SegmentTiming) {
 		segmentTiming.text = segment.text
 		if segmentTiming.active {
 			b.renderSegment(segment)
-			segmentTiming.text = b.writer.String()
+			segmentTiming.text, b.length = b.writer.String()
 			b.writer.Reset()
 		}
 		segmentTiming.duration = time.Since(start)

src/engine/engine.go

@@ -19,7 +19,9 @@ type Engine struct {
 	Plain        bool
 
 	console           strings.Builder
+	currentLineLength int
 	rprompt           string
+	rpromptLength     int
 }
 
 func (e *Engine) write(text string) {
@@ -38,12 +40,11 @@ func (e *Engine) string() string {
 }
 
 func (e *Engine) canWriteRPrompt() bool {
-	prompt := e.string()
 	consoleWidth, err := e.Env.TerminalWidth()
 	if err != nil || consoleWidth == 0 {
 		return true
 	}
-	promptWidth := e.Ansi.LenWithoutANSI(prompt)
+	promptWidth := e.currentLineLength
 	availableSpace := consoleWidth - promptWidth
 	// spanning multiple lines
 	if availableSpace < 0 {
@@ -51,7 +52,7 @@ func (e *Engine) canWriteRPrompt() bool {
 		availableSpace = consoleWidth - overflow
 	}
 	promptBreathingRoom := 30
-	canWrite := (availableSpace - e.Ansi.LenWithoutANSI(e.rprompt)) >= promptBreathingRoom
+	canWrite := (availableSpace - e.rpromptLength) >= promptBreathingRoom
 	return canWrite
 }
 
@@ -66,7 +67,6 @@ func (e *Engine) Render() string {
 	if e.Config.FinalSpace {
 		e.write(" ")
 	}
-
 	if !e.Config.OSC99 {
 		return e.print()
 	}
@@ -75,6 +75,11 @@ func (e *Engine) Render() string {
 	return e.print()
 }
 
+func (e *Engine) newline() {
+	e.write("\n")
+	e.currentLineLength = 0
+}
+
 func (e *Engine) renderBlock(block *Block) {
 	// when in bash, for rprompt blocks we need to write plain
 	// and wrap in escaped mode or the prompt will not render correctly
@@ -88,14 +93,14 @@ func (e *Engine) renderBlock(block *Block) {
 		return
 	}
 	if block.Newline {
-		e.write("\n")
+		e.newline()
 	}
 	switch block.Type {
 	// This is deprecated but leave if to not break current configs
 	// It is encouraged to used "newline": true on block level
 	// rather than the standalone the linebreak block
 	case LineBreak:
-		e.write("\n")
+		e.newline()
 	case Prompt:
 		if block.VerticalOffset != 0 {
 			e.writeANSI(e.Ansi.ChangeLine(block.VerticalOffset))
@@ -103,18 +108,22 @@ func (e *Engine) renderBlock(block *Block) {
 		switch block.Alignment {
 		case Right:
 			e.writeANSI(e.Ansi.CarriageForward())
-			blockText := block.renderSegments()
+			text, length := block.renderSegments()
-			e.writeANSI(e.Ansi.GetCursorForRightWrite(blockText, block.HorizontalOffset))
+			e.currentLineLength += length
-			e.write(blockText)
+			e.writeANSI(e.Ansi.GetCursorForRightWrite(length, block.HorizontalOffset))
+			e.write(text)
 		case Left:
-			e.write(block.renderSegments())
+			text, length := block.renderSegments()
+			e.currentLineLength += length
+			e.write(text)
 		}
 	case RPrompt:
-		blockText := block.renderSegments()
+		text, length := block.renderSegments()
+		e.rpromptLength = length
 		if e.Env.Shell() == bash {
-			blockText = e.Ansi.FormatText(blockText)
+			text = e.Ansi.FormatText(text)
 		}
-		e.rprompt = blockText
+		e.rprompt = text
 	}
 	// Due to a bug in Powershell, the end of the line needs to be cleared.
 	// If this doesn't happen, the portion after the prompt gets colored in the background
@@ -183,7 +192,7 @@ func (e *Engine) print() string {
 		}
 		e.write(e.Ansi.SaveCursorPosition())
 		e.write(e.Ansi.CarriageForward())
-		e.write(e.Ansi.GetCursorForRightWrite(e.rprompt, 0))
+		e.write(e.Ansi.GetCursorForRightWrite(e.rpromptLength, 0))
 		e.write(e.rprompt)
 		e.write(e.Ansi.RestoreCursorPosition())
 	}
@@ -217,14 +226,15 @@ func (e *Engine) RenderTooltip(tip string) string {
 	switch e.Env.Shell() {
 	case zsh, winCMD:
 		block.init(e.Env, e.Writer, e.Ansi)
-		return block.renderSegments()
+		text, _ := block.renderSegments()
+		return text
 	case pwsh, powershell5:
 		block.initPlain(e.Env, e.Config)
-		tooltipText := block.renderSegments()
+		text, length := block.renderSegments()
 		e.write(e.Ansi.ClearAfter())
 		e.write(e.Ansi.CarriageForward())
-		e.write(e.Ansi.GetCursorForRightWrite(tooltipText, 0))
+		e.write(e.Ansi.GetCursorForRightWrite(length, 0))
-		e.write(tooltipText)
+		e.write(text)
 		return e.string()
 	}
 	return ""
@@ -251,11 +261,13 @@ func (e *Engine) RenderTransientPrompt() string {
 	switch e.Env.Shell() {
 	case zsh:
 		// escape double quotes contained in the prompt
-		prompt := fmt.Sprintf("PS1=\"%s\"", strings.ReplaceAll(e.Writer.String(), "\"", "\"\""))
+		str, _ := e.Writer.String()
+		prompt := fmt.Sprintf("PS1=\"%s\"", strings.ReplaceAll(str, "\"", "\"\""))
 		prompt += "\nRPROMPT=\"\""
 		return prompt
 	case pwsh, powershell5, winCMD:
-		return e.Writer.String()
+		str, _ := e.Writer.String()
+		return str
 	}
 	return ""
 }
@@ -278,5 +290,7 @@ func (e *Engine) RenderRPrompt() string {
 	if !block.enabled() {
 		return ""
 	}
-	return block.renderSegments()
+	text, length := block.renderSegments()
+	e.rpromptLength = length
+	return text
 }

src/engine/engine_test.go

@@ -8,7 +8,6 @@ import (
 	"oh-my-posh/mock"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -26,8 +25,8 @@ func TestCanWriteRPrompt(t *testing.T) {
 		{Case: "Width Error", Expected: true, TerminalWidthError: errors.New("burp")},
 		{Case: "Terminal > Prompt enabled", Expected: true, TerminalWidth: 200, PromptLength: 100, RPromptLength: 10},
 		{Case: "Terminal > Prompt enabled edge", Expected: true, TerminalWidth: 200, PromptLength: 100, RPromptLength: 70},
-		{Case: "Terminal > Prompt disabled no breathing", Expected: false, TerminalWidth: 200, PromptLength: 100, RPromptLength: 71},
 		{Case: "Prompt > Terminal enabled", Expected: true, TerminalWidth: 200, PromptLength: 300, RPromptLength: 70},
+		{Case: "Terminal > Prompt disabled no breathing", Expected: false, TerminalWidth: 200, PromptLength: 100, RPromptLength: 71},
 		{Case: "Prompt > Terminal disabled no breathing", Expected: false, TerminalWidth: 200, PromptLength: 300, RPromptLength: 80},
 		{Case: "Prompt > Terminal disabled no room", Expected: true, TerminalWidth: 200, PromptLength: 400, RPromptLength: 80},
 	}
@@ -35,14 +34,11 @@ func TestCanWriteRPrompt(t *testing.T) {
 	for _, tc := range cases {
 		env := new(mock.MockedEnvironment)
 		env.On("TerminalWidth").Return(tc.TerminalWidth, tc.TerminalWidthError)
-		ansi := &color.Ansi{}
-		ansi.Init(plain)
 		engine := &Engine{
 			Env: env,
-			Ansi: ansi,
 		}
-		engine.rprompt = strings.Repeat("x", tc.RPromptLength)
+		engine.rpromptLength = tc.RPromptLength
-		engine.console.WriteString(strings.Repeat("x", tc.PromptLength))
+		engine.currentLineLength = tc.PromptLength
 		got := engine.canWriteRPrompt()
 		assert.Equal(t, tc.Expected, got, tc.Case)
 	}

src/engine/image.go

@@ -67,6 +67,8 @@ const (
 	lineChange          = "linechange"
 	consoleTitle        = "title"
 	link                = "link"
+
+	ansiRegex = "[\u001B\u009B][[\\]()#;?]*(?:(?:(?:[a-zA-Z\\d]*(?:;[a-zA-Z\\d]*)*)?\u0007)|(?:(?:\\d{1,4}(?:;\\d{0,4})*)?[\\dA-PRZcf-ntqry=><~]))"
 )
 
 //go:embed font/Hack-Nerd-Bold.ttf
@@ -187,7 +189,7 @@ func (ir *ImageRenderer) Init(config string) {
 		osc99:               `^(?P<STR>\x1b\]9;9;(.+)\x1b\\)`,
 		lineChange:          `^(?P<STR>\x1b\[(\d)[FB])`,
 		consoleTitle:        `^(?P<STR>\x1b\]0;(.+)\007)`,
-		link:                `^(?P<STR>\x1b]8;;file:\/\/(.+)\x1b\\(?P<URL>.+)\x1b]8;;\x1b\\)`,
+		link:                `^(?P<STR>\x1b]8;;(file|https)(.+)\x1b\\(?P<URL>.+)\x1b]8;;\x1b\\)`,
 	}
 }
 
@@ -247,13 +249,35 @@ func (ir *ImageRenderer) runeAdditionalWidth(r rune) int {
 	return 0
 }
 
+func (ir *ImageRenderer) lenWithoutANSI(text string) int {
+	if len(text) == 0 {
+		return 0
+	}
+	// replace hyperlinks(file/http/https)
+	regexStr := ir.ansiSequenceRegexMap[link]
+	matches := regex.FindAllNamedRegexMatch(regexStr, text)
+	for _, match := range matches {
+		text = strings.ReplaceAll(text, match[str], match[url])
+	}
+	// replace console title
+	regexStr = ir.ansiSequenceRegexMap[consoleTitle]
+	matches = regex.FindAllNamedRegexMatch(regexStr, text)
+	for _, match := range matches {
+		text = strings.ReplaceAll(text, match[str], "")
+	}
+	stripped := regex.ReplaceAllString(ansiRegex, text, "")
+	runeText := []rune(stripped)
+	length := len(runeText)
+	for _, rune := range runeText {
+		length += ir.runeAdditionalWidth(rune)
+	}
+	return length
+}
+
 func (ir *ImageRenderer) calculateWidth() int {
 	longest := 0
 	for _, line := range strings.Split(ir.AnsiString, "\n") {
-		length := ir.Ansi.LenWithoutANSI(line)
+		length := ir.lenWithoutANSI(line)
-		for _, char := range line {
-			length += ir.runeAdditionalWidth(char)
-		}
 		if length > longest {
 			longest = length
 		}

src/engine/image_test.go

@@ -33,9 +33,7 @@ func TestStringImageFileWithText(t *testing.T) {
 }
 
 func TestStringImageFileWithANSI(t *testing.T) {
-	prompt := `\uE0B0 oh-my-posh
+	prompt := ` jan  `
-	 \uE0B0 main ≡  ~4 -8 ?7 \uE0B0
-	 `
 	err := runImageTest(prompt)
 	assert.NoError(t, err)
 }