fix(cli): route stream output through injected stdout writer
Some checks are pending
Security Scan / security (push) Waiting to run
Some checks are pending
Security Scan / security (push) Waiting to run
Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
parent
050ee5bd8f
commit
cdae3a9ac5
3 changed files with 20 additions and 6 deletions
|
|
@ -34,7 +34,8 @@ When word-wrap is enabled, the stream tracks the current column position and ins
|
||||||
|
|
||||||
## Custom Output Writer
|
## Custom Output Writer
|
||||||
|
|
||||||
By default, streams write to `os.Stdout`. Redirect to any `io.Writer`:
|
By default, streams write to the CLI stdout writer (`stdoutWriter()`), so tests can
|
||||||
|
redirect output via `cli.SetStdout` and other callers can provide any `io.Writer`:
|
||||||
|
|
||||||
```go
|
```go
|
||||||
var buf strings.Builder
|
var buf strings.Builder
|
||||||
|
|
@ -77,7 +78,7 @@ stream.Done()
|
||||||
| Option | Description |
|
| Option | Description |
|
||||||
|--------|-------------|
|
|--------|-------------|
|
||||||
| `WithWordWrap(cols)` | Set the word-wrap column width |
|
| `WithWordWrap(cols)` | Set the word-wrap column width |
|
||||||
| `WithStreamOutput(w)` | Set the output writer (default: `os.Stdout`) |
|
| `WithStreamOutput(w)` | Set the output writer (default: `stdoutWriter()`) |
|
||||||
|
|
||||||
## Example: LLM Token Streaming
|
## Example: LLM Token Streaming
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,6 @@ package cli
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"os"
|
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
|
|
@ -12,7 +11,8 @@ import (
|
||||||
|
|
||||||
// StreamOption configures a Stream.
|
// StreamOption configures a Stream.
|
||||||
//
|
//
|
||||||
// stream := cli.NewStream(cli.WithWordWrap(80), cli.WithStreamOutput(os.Stdout))
|
// stream := cli.NewStream(cli.WithWordWrap(80))
|
||||||
|
// stream.Wait()
|
||||||
type StreamOption func(*Stream)
|
type StreamOption func(*Stream)
|
||||||
|
|
||||||
// WithWordWrap sets the word-wrap column width.
|
// WithWordWrap sets the word-wrap column width.
|
||||||
|
|
@ -20,7 +20,7 @@ func WithWordWrap(cols int) StreamOption {
|
||||||
return func(s *Stream) { s.wrap = cols }
|
return func(s *Stream) { s.wrap = cols }
|
||||||
}
|
}
|
||||||
|
|
||||||
// WithStreamOutput sets the output writer (default: os.Stdout).
|
// WithStreamOutput sets the output writer (default: stdoutWriter()).
|
||||||
func WithStreamOutput(w io.Writer) StreamOption {
|
func WithStreamOutput(w io.Writer) StreamOption {
|
||||||
return func(s *Stream) { s.out = w }
|
return func(s *Stream) { s.out = w }
|
||||||
}
|
}
|
||||||
|
|
@ -48,7 +48,7 @@ type Stream struct {
|
||||||
// NewStream creates a streaming text renderer.
|
// NewStream creates a streaming text renderer.
|
||||||
func NewStream(opts ...StreamOption) *Stream {
|
func NewStream(opts ...StreamOption) *Stream {
|
||||||
s := &Stream{
|
s := &Stream{
|
||||||
out: os.Stdout,
|
out: stdoutWriter(),
|
||||||
done: make(chan struct{}),
|
done: make(chan struct{}),
|
||||||
}
|
}
|
||||||
for _, opt := range opts {
|
for _, opt := range opts {
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,19 @@ import (
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestStream_Good(t *testing.T) {
|
func TestStream_Good(t *testing.T) {
|
||||||
|
t.Run("uses injected stdout by default", func(t *testing.T) {
|
||||||
|
var buf bytes.Buffer
|
||||||
|
SetStdout(&buf)
|
||||||
|
defer SetStdout(nil)
|
||||||
|
s := NewStream()
|
||||||
|
|
||||||
|
s.Write("hello")
|
||||||
|
s.Done()
|
||||||
|
s.Wait()
|
||||||
|
|
||||||
|
assert.Equal(t, "hello\n", buf.String())
|
||||||
|
})
|
||||||
|
|
||||||
t.Run("basic write", func(t *testing.T) {
|
t.Run("basic write", func(t *testing.T) {
|
||||||
var buf bytes.Buffer
|
var buf bytes.Buffer
|
||||||
s := NewStream(WithStreamOutput(&buf))
|
s := NewStream(WithStreamOutput(&buf))
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue