feat: add Wait to wait for expected output (#257)

* add Match and MatchAny commands

* allow changing Match value

* switch to MatchLine and MatchScreen

* remove leftover code

* add wait and regex to lexer

* add parseWait

* Add wait option defaults

* add execute for wait

* allow setting WAitTimeout and WaitPattern

* remove MATCH_LINE and MATCH_SCREEN

* add description to Buffer method

* add wait to parser test and fix setting parsing

* update TestCommand

* improve error output for Wait timeout

* don't require regex (fall back to WaitPattern)

* update signatures with errors

* set default wait timeout to 15s

* Update testing.go

Co-authored-by: Christian Rocha <[email protected]>

---------

Co-authored-by: Christian Rocha <[email protected]>

command.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -76,6 +77,7 @@ var CommandFuncs = map[parser.CommandType]CommandFunc{
 	token.COPY:       ExecuteCopy,
 	token.PASTE:      ExecutePaste,
 	token.ENV:        ExecuteEnv,
+	token.WAIT:       ExecuteWait,
 }
 
 // ExecuteNoop is a no-op command that does nothing.
@@ -111,6 +113,71 @@ func ExecuteKey(k input.Key) CommandFunc {
 	}
 }
 
+// WaitTick is the amount of time to wait between checking for a match.
+const WaitTick = 10 * time.Millisecond
+
+// ExecuteWait is a CommandFunc that waits for a regex match for the given amount of time.
+func ExecuteWait(c parser.Command, v *VHS) error {
+	scope, rxStr, ok := strings.Cut(c.Args, " ")
+	rx := v.Options.WaitPattern
+	if ok {
+		// This is validated on parse so using MustCompile reduces noise.
+		rx = regexp.MustCompile(rxStr)
+	}
+
+	timeout := v.Options.WaitTimeout
+	if c.Options != "" {
+		t, err := time.ParseDuration(c.Options)
+		if err != nil {
+			// Shouldn't be possible due to parse validation.
+			return fmt.Errorf("failed to parse duration: %w", err)
+		}
+		timeout = t
+	}
+
+	checkT := time.NewTicker(WaitTick)
+	defer checkT.Stop()
+	timeoutT := time.NewTimer(timeout)
+	defer timeoutT.Stop()
+
+	for {
+		var last string
+		switch scope {
+		case "Line":
+			line, err := v.CurrentLine()
+			if err != nil {
+				return fmt.Errorf("failed to get current line: %w", err)
+			}
+			last = line
+
+			if rx.MatchString(line) {
+				return nil
+			}
+		case "Screen":
+			lines, err := v.Buffer()
+			if err != nil {
+				return fmt.Errorf("failed to get buffer: %w", err)
+			}
+			last = strings.Join(lines, "\n")
+
+			if rx.MatchString(last) {
+				return nil
+			}
+		default:
+			// Should be impossible due to parse validation, but we don't want to
+			// hang if it does happen due to a bug.
+			return fmt.Errorf("invalid scope %q", scope)
+		}
+
+		select {
+		case <-checkT.C:
+			continue
+		case <-timeoutT.C:
+			return fmt.Errorf("timeout waiting for %q to match %s; last value was: %s", c.Args, rx.String(), last)
+		}
+	}
+}
+
 // ExecuteCtrl is a CommandFunc that presses the argument keys and/or modifiers
 // with the ctrl key held down on the running instance of vhs.
 func ExecuteCtrl(c parser.Command, v *VHS) error {
@@ -371,6 +438,8 @@ var Settings = map[string]CommandFunc{
 	"WindowBar":     ExecuteSetWindowBar,
 	"WindowBarSize": ExecuteSetWindowBarSize,
 	"BorderRadius":  ExecuteSetBorderRadius,
+	"WaitPattern":   ExecuteSetWaitPattern,
+	"WaitTimeout":   ExecuteSetWaitTimeout,
 	"CursorBlink":   ExecuteSetCursorBlink,
 }
 
@@ -521,6 +590,26 @@ func ExecuteSetTypingSpeed(c parser.Command, v *VHS) error {
 	return nil
 }
 
+// ExecuteSetWaitTimeout applies the default wait timeout on the vhs.
+func ExecuteSetWaitTimeout(c parser.Command, v *VHS) error {
+	waitTimeout, err := time.ParseDuration(c.Args)
+	if err != nil {
+		return fmt.Errorf("failed to parse wait timeout: %w", err)
+	}
+	v.Options.WaitTimeout = waitTimeout
+	return nil
+}
+
+// ExecuteSetWaitPattern applies the default wait pattern on the vhs.
+func ExecuteSetWaitPattern(c parser.Command, v *VHS) error {
+	rx, err := regexp.Compile(c.Args)
+	if err != nil {
+		return fmt.Errorf("failed to compile regexp: %w", err)
+	}
+	v.Options.WaitPattern = rx
+	return nil
+}
+
 // ExecuteSetPadding applies the padding on the vhs.
 func ExecuteSetPadding(c parser.Command, v *VHS) error {
 	padding, err := strconv.Atoi(c.Args)

command_test.go

@@ -8,12 +8,12 @@ import (
 )
 
 func TestCommand(t *testing.T) {
-	const numberOfCommands = 28
+	const numberOfCommands = 29
 	if len(parser.CommandTypes) != numberOfCommands {
 		t.Errorf("Expected %d commands, got %d", numberOfCommands, len(parser.CommandTypes))
 	}
 
-	const numberOfCommandFuncs = 28
+	const numberOfCommandFuncs = 29
 	if len(CommandFuncs) != numberOfCommandFuncs {
 		t.Errorf("Expected %d commands, got %d", numberOfCommandFuncs, len(CommandFuncs))
 	}

lexer/lexer.go

@@ -67,6 +67,10 @@ func (l *Lexer) NextToken() token.Token {
 		tok.Type = token.STRING
 		tok.Literal = l.readString('"')
 		l.readChar()
+	case '/':
+		tok.Type = token.REGEX
+		tok.Literal = l.readString('/')
+		l.readChar()
 	default:
 		if isDigit(l.ch) || (isDot(l.ch) && isDigit(l.peekChar())) {
 			tok.Literal = l.readNumber()

lexer/lexer_test.go

@@ -24,7 +24,8 @@ Ctrl+C
 Enter
 Sleep .1
 Sleep 100ms
-Sleep 2`
+Sleep 2
+Wait+Screen@1m /foobar/`
 
 	tests := []struct {
 		expectedType    token.Type
@@ -71,6 +72,13 @@ Sleep 2`
 		{token.MILLISECONDS, "ms"},
 		{token.SLEEP, "Sleep"},
 		{token.NUMBER, "2"},
+		{token.WAIT, "Wait"},
+		{token.PLUS, "+"},
+		{token.STRING, "Screen"},
+		{token.AT, "@"},
+		{token.NUMBER, "1"},
+		{token.MINUTES, "m"},
+		{token.REGEX, "foobar"},
 	}
 
 	l := New(input)

man.go

@@ -43,6 +43,7 @@ The following is a list of all possible commands in VHS:
 * %PageDown% [repeat]
 * %Hide%
 * %Show%
+* %Wait%[+Screen][@<timeout>] /<regexp>/
 * %Escape%
 * %Alt%+<key>
 * %Space% [repeat]
@@ -72,6 +73,8 @@ The following is a list of all possible setting commands in VHS:
 * Set %Padding% <number>
 * Set %Framerate% <number>
 * Set %PlaybackSpeed% <float>
+* Set %WaitTimeout% <time>
+* Set %WaitPattern% <regexp>
 `
 	manBugs = "See GitHub Issues: <https://github.com/charmbracelet/vhs/issues>"
 

parser/parser.go

@@ -4,8 +4,10 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/charmbracelet/vhs/lexer"
 	"github.com/charmbracelet/vhs/token"
@@ -47,6 +49,7 @@ var CommandTypes = []CommandType{ //nolint: deadcode
 	token.TAB,
 	token.TYPE,
 	token.UP,
+	token.WAIT,
 	token.SOURCE,
 	token.SCREENSHOT,
 	token.COPY,
@@ -55,13 +58,7 @@ var CommandTypes = []CommandType{ //nolint: deadcode
 }
 
 // String returns the string representation of the command.
-func (c CommandType) String() string {
-	if len(c) < 1 {
-		return ""
-	}
-	s := string(c)
-	return string(s[0]) + strings.ToLower(s[1:])
-}
+func (c CommandType) String() string { return token.ToCamel(string(c)) }
 
 // Command represents a command with options and arguments.
 type Command struct {
@@ -170,6 +167,8 @@ func (p *Parser) parseCommand() Command {
 		return p.parseRequire()
 	case token.SHOW:
 		return p.parseShow()
+	case token.WAIT:
+		return p.parseWait()
 	case token.SOURCE:
 		return p.parseSource()
 	case token.SCREENSHOT:
@@ -186,6 +185,45 @@ func (p *Parser) parseCommand() Command {
 	}
 }
 
+func (p *Parser) parseWait() Command {
+	cmd := Command{Type: token.WAIT}
+
+	if p.peek.Type == token.PLUS {
+		p.nextToken()
+		if p.peek.Type != token.STRING || (p.peek.Literal != "Line" && p.peek.Literal != "Screen") {
+			p.errors = append(p.errors, NewError(p.peek, "Wait+ expects Line or Screen"))
+			return cmd
+		}
+		cmd.Args = p.peek.Literal
+		p.nextToken()
+	} else {
+		cmd.Args = "Line"
+	}
+
+	cmd.Options = p.parseSpeed()
+	if cmd.Options != "" {
+		dur, _ := time.ParseDuration(cmd.Options)
+		if dur <= 0 {
+			p.errors = append(p.errors, NewError(p.peek, "Wait expects positive duration"))
+			return cmd
+		}
+	}
+
+	if p.peek.Type != token.REGEX {
+		// fallback to default
+		return cmd
+	}
+	p.nextToken()
+	if _, err := regexp.Compile(p.cur.Literal); err != nil {
+		p.errors = append(p.errors, NewError(p.cur, fmt.Sprintf("Invalid regular expression '%s': %v", p.cur.Literal, err)))
+		return cmd
+	}
+
+	cmd.Args += " " + p.cur.Literal
+
+	return cmd
+}
+
 // parseSpeed parses a typing speed indication.
 //
 // i.e. @<time>
@@ -227,10 +265,11 @@ func (p *Parser) parseTime() string {
 		p.nextToken()
 	} else {
 		p.errors = append(p.errors, NewError(p.cur, "Expected time after "+p.cur.Literal))
+		return ""
 	}
 
 	// Allow TypingSpeed to have bare units (e.g. 50ms, 100ms)
-	if p.peek.Type == token.MILLISECONDS || p.peek.Type == token.SECONDS {
+	if p.peek.Type == token.MILLISECONDS || p.peek.Type == token.SECONDS || p.peek.Type == token.MINUTES {
 		t += p.peek.Literal
 		p.nextToken()
 	} else {
@@ -393,6 +432,15 @@ func (p *Parser) parseSet() Command {
 	p.nextToken()
 
 	switch p.cur.Type {
+	case token.WAIT_TIMEOUT:
+		cmd.Args = p.parseTime()
+	case token.WAIT_PATTERN:
+		cmd.Args = p.peek.Literal
+		_, err := regexp.Compile(p.peek.Literal)
+		if err != nil {
+			p.errors = append(p.errors, NewError(p.peek, "Invalid regexp pattern: "+p.peek.Literal))
+		}
+		p.nextToken()
 	case token.LOOP_OFFSET:
 		cmd.Args = p.peek.Literal
 		p.nextToken()

parser/parser_test.go

@@ -12,6 +12,8 @@ import (
 func TestParser(t *testing.T) {
 	input := `
 Set TypingSpeed 100ms
+Set WaitTimeout 1m
+Set WaitPattern /foo/
 Type "echo 'Hello, World!'"
 Enter
 [email protected] 5
@@ -28,10 +30,15 @@ Ctrl+C
 Ctrl+L
 Alt+.
 Sleep 100ms
-Sleep 3`
+Sleep 3
+Wait
+Wait+Screen
+Wait@100ms /foobar/`
 
 	expected := []Command{
 		{Type: token.SET, Options: "TypingSpeed", Args: "100ms"},
+		{Type: token.SET, Options: "WaitTimeout", Args: "1m"},
+		{Type: token.SET, Options: "WaitPattern", Args: "foo"},
 		{Type: token.TYPE, Options: "", Args: "echo 'Hello, World!'"},
 		{Type: token.ENTER, Options: "", Args: "1"},
 		{Type: token.BACKSPACE, Options: "0.1s", Args: "5"},
@@ -49,6 +56,9 @@ Sleep 3`
 		{Type: token.ALT, Options: "", Args: "."},
 		{Type: token.SLEEP, Args: "100ms"},
 		{Type: token.SLEEP, Args: "3s"},
+		{Type: token.WAIT, Args: "Line"},
+		{Type: token.WAIT, Args: "Screen"},
+		{Type: token.WAIT, Options: "100ms", Args: "Line foobar"},
 	}
 
 	l := lexer.New(input)
@@ -57,7 +67,7 @@ Sleep 3`
 	cmds := p.Parse()
 
 	if len(cmds) != len(expected) {
-		t.Fatalf("Expected %d commands, got %d", len(expected), len(cmds))
+		t.Fatalf("Expected %d commands, got %d; %v", len(expected), len(cmds), cmds)
 	}
 
 	for i, cmd := range cmds {

syntax.go

@@ -26,6 +26,8 @@ func Highlight(c parser.Command, faint bool) string {
 	}
 
 	switch c.Type {
+	case token.REGEX:
+		argsStyle = StringStyle
 	case token.SET:
 		optionsStyle = KeywordStyle
 		if isNumber(c.Args) {