github.com/rentongzhang/docker@v1.8.2-rc1/pkg/streamformatter/streamformatter.go (about)

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