github.com/lmars/docker@v1.6.0-rc2/utils/streamformatter.go (about)

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