github.com/shiroyuki/docker@v1.9.0/pkg/streamformatter/streamformatter.go (about)

     1  // Package streamformatter provides helper functions to format a stream.
     2  package streamformatter
     3  
     4  import (
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  
     9  	"github.com/docker/docker/pkg/jsonmessage"
    10  )
    11  
    12  // StreamFormatter formats a stream, optionally using JSON.
    13  type StreamFormatter struct {
    14  	json bool
    15  }
    16  
    17  // NewStreamFormatter returns a simple StreamFormatter
    18  func NewStreamFormatter() *StreamFormatter {
    19  	return &StreamFormatter{}
    20  }
    21  
    22  // NewJSONStreamFormatter returns a StreamFormatter configured to stream json
    23  func NewJSONStreamFormatter() *StreamFormatter {
    24  	return &StreamFormatter{true}
    25  }
    26  
    27  const streamNewline = "\r\n"
    28  
    29  var streamNewlineBytes = []byte(streamNewline)
    30  
    31  // FormatStream formats the specified stream.
    32  func (sf *StreamFormatter) FormatStream(str string) []byte {
    33  	if sf.json {
    34  		b, err := json.Marshal(&jsonmessage.JSONMessage{Stream: str})
    35  		if err != nil {
    36  			return sf.FormatError(err)
    37  		}
    38  		return append(b, streamNewlineBytes...)
    39  	}
    40  	return []byte(str + "\r")
    41  }
    42  
    43  // FormatStatus formats the specified objects according to the specified format (and id).
    44  func (sf *StreamFormatter) FormatStatus(id, format string, a ...interface{}) []byte {
    45  	str := fmt.Sprintf(format, a...)
    46  	if sf.json {
    47  		b, err := json.Marshal(&jsonmessage.JSONMessage{ID: id, Status: str})
    48  		if err != nil {
    49  			return sf.FormatError(err)
    50  		}
    51  		return append(b, streamNewlineBytes...)
    52  	}
    53  	return []byte(str + streamNewline)
    54  }
    55  
    56  // FormatError formats the specifed error.
    57  func (sf *StreamFormatter) FormatError(err error) []byte {
    58  	if sf.json {
    59  		jsonError, ok := err.(*jsonmessage.JSONError)
    60  		if !ok {
    61  			jsonError = &jsonmessage.JSONError{Message: err.Error()}
    62  		}
    63  		if b, err := json.Marshal(&jsonmessage.JSONMessage{Error: jsonError, ErrorMessage: err.Error()}); err == nil {
    64  			return append(b, streamNewlineBytes...)
    65  		}
    66  		return []byte("{\"error\":\"format error\"}" + streamNewline)
    67  	}
    68  	return []byte("Error: " + err.Error() + streamNewline)
    69  }
    70  
    71  // FormatProgress formats the progress information for a specified action.
    72  func (sf *StreamFormatter) FormatProgress(id, action string, progress *jsonmessage.JSONProgress) []byte {
    73  	if progress == nil {
    74  		progress = &jsonmessage.JSONProgress{}
    75  	}
    76  	if sf.json {
    77  		b, err := json.Marshal(&jsonmessage.JSONMessage{
    78  			Status:          action,
    79  			ProgressMessage: progress.String(),
    80  			Progress:        progress,
    81  			ID:              id,
    82  		})
    83  		if err != nil {
    84  			return nil
    85  		}
    86  		return b
    87  	}
    88  	endl := "\r"
    89  	if progress.String() == "" {
    90  		endl += "\n"
    91  	}
    92  	return []byte(action + " " + progress.String() + endl)
    93  }
    94  
    95  // StdoutFormatter is a streamFormatter that writes to the standard output.
    96  type StdoutFormatter struct {
    97  	io.Writer
    98  	*StreamFormatter
    99  }
   100  
   101  func (sf *StdoutFormatter) Write(buf []byte) (int, error) {
   102  	formattedBuf := sf.StreamFormatter.FormatStream(string(buf))
   103  	n, err := sf.Writer.Write(formattedBuf)
   104  	if n != len(formattedBuf) {
   105  		return n, io.ErrShortWrite
   106  	}
   107  	return len(buf), err
   108  }
   109  
   110  // StderrFormatter is a streamFormatter that writes to the standard error.
   111  type StderrFormatter struct {
   112  	io.Writer
   113  	*StreamFormatter
   114  }
   115  
   116  func (sf *StderrFormatter) Write(buf []byte) (int, error) {
   117  	formattedBuf := sf.StreamFormatter.FormatStream("\033[91m" + string(buf) + "\033[0m")
   118  	n, err := sf.Writer.Write(formattedBuf)
   119  	if n != len(formattedBuf) {
   120  		return n, io.ErrShortWrite
   121  	}
   122  	return len(buf), err
   123  }