github.com/telepresenceio/telepresence/v2@v2.20.0-pro.6.0.20240517030216-236ea954e789/pkg/dos/stdio.go (about)

     1  package dos
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"os"
     7  )
     8  
     9  type Stdio interface {
    10  	InOrStdin() io.Reader
    11  	OutOrStdout() io.Writer
    12  	ErrOrStderr() io.Writer
    13  }
    14  
    15  // WithStdio returns a new context that is parented by the given context, that will return the values
    16  // from the given Stdio when used as an argument to in a call to Stdin, Stdout, or Stderr.
    17  func WithStdio(ctx context.Context, io Stdio) context.Context {
    18  	ctx = WithStdin(ctx, io.InOrStdin())
    19  	ctx = WithStdout(ctx, io.OutOrStdout())
    20  	return WithStderr(ctx, io.ErrOrStderr())
    21  }
    22  
    23  type outKey struct{}
    24  
    25  // Stdout returns the context's stdout io.Writer. If no such writer has been defined, it returns os.Stdout.
    26  func Stdout(ctx context.Context) io.Writer {
    27  	if w, ok := ctx.Value(outKey{}).(io.Writer); ok {
    28  		return w
    29  	}
    30  	return os.Stdout
    31  }
    32  
    33  // WithStdout returns a new context that is parented by the given context, that will return the given writer
    34  // when used as an argument to in a call to Stdout.
    35  func WithStdout(ctx context.Context, w io.Writer) context.Context {
    36  	return context.WithValue(ctx, outKey{}, w)
    37  }
    38  
    39  type errKey struct{}
    40  
    41  // Stderr returns the context's stdout io.Writer. If no such writer has been defined, it returns os.Stderr.
    42  func Stderr(ctx context.Context) io.Writer {
    43  	if w, ok := ctx.Value(errKey{}).(io.Writer); ok {
    44  		return w
    45  	}
    46  	return os.Stderr
    47  }
    48  
    49  // WithStderr returns a new context that is parented by the given context, that will return the given writer
    50  // when used as an argument to in a call to Stderr.
    51  func WithStderr(ctx context.Context, w io.Writer) context.Context {
    52  	return context.WithValue(ctx, errKey{}, w)
    53  }
    54  
    55  type inKey struct{}
    56  
    57  // Stdin returns the context's stdin io.Reader. If no such reader has been defined, it returns os.Stdin.
    58  func Stdin(ctx context.Context) io.Reader {
    59  	if r, ok := ctx.Value(inKey{}).(io.Reader); ok {
    60  		return r
    61  	}
    62  	return os.Stdin
    63  }
    64  
    65  // WithStdin returns a new context that is parented by the given context, that will return the given reader
    66  // when used as an argument to in a call to Stdin.
    67  func WithStdin(ctx context.Context, w io.Reader) context.Context {
    68  	return context.WithValue(ctx, inKey{}, w)
    69  }