github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/pkg/streamformatter/streamformatter.go (about)

     1  // Package streamformatter provides helper functions to format a stream.
     2  package streamformatter // import "github.com/demonoid81/moby/pkg/streamformatter"
     3  
     4  import (
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  
     9  	"github.com/demonoid81/moby/pkg/jsonmessage"
    10  	"github.com/demonoid81/moby/pkg/progress"
    11  )
    12  
    13  const streamNewline = "\r\n"
    14  
    15  type jsonProgressFormatter struct{}
    16  
    17  func appendNewline(source []byte) []byte {
    18  	return append(source, []byte(streamNewline)...)
    19  }
    20  
    21  // FormatStatus formats the specified objects according to the specified format (and id).
    22  func FormatStatus(id, format string, a ...interface{}) []byte {
    23  	str := fmt.Sprintf(format, a...)
    24  	b, err := json.Marshal(&jsonmessage.JSONMessage{ID: id, Status: str})
    25  	if err != nil {
    26  		return FormatError(err)
    27  	}
    28  	return appendNewline(b)
    29  }
    30  
    31  // FormatError formats the error as a JSON object
    32  func FormatError(err error) []byte {
    33  	jsonError, ok := err.(*jsonmessage.JSONError)
    34  	if !ok {
    35  		jsonError = &jsonmessage.JSONError{Message: err.Error()}
    36  	}
    37  	if b, err := json.Marshal(&jsonmessage.JSONMessage{Error: jsonError, ErrorMessage: err.Error()}); err == nil {
    38  		return appendNewline(b)
    39  	}
    40  	return []byte(`{"error":"format error"}` + streamNewline)
    41  }
    42  
    43  func (sf *jsonProgressFormatter) formatStatus(id, format string, a ...interface{}) []byte {
    44  	return FormatStatus(id, format, a...)
    45  }
    46  
    47  // formatProgress formats the progress information for a specified action.
    48  func (sf *jsonProgressFormatter) formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte {
    49  	if progress == nil {
    50  		progress = &jsonmessage.JSONProgress{}
    51  	}
    52  	var auxJSON *json.RawMessage
    53  	if aux != nil {
    54  		auxJSONBytes, err := json.Marshal(aux)
    55  		if err != nil {
    56  			return nil
    57  		}
    58  		auxJSON = new(json.RawMessage)
    59  		*auxJSON = auxJSONBytes
    60  	}
    61  	b, err := json.Marshal(&jsonmessage.JSONMessage{
    62  		Status:          action,
    63  		ProgressMessage: progress.String(),
    64  		Progress:        progress,
    65  		ID:              id,
    66  		Aux:             auxJSON,
    67  	})
    68  	if err != nil {
    69  		return nil
    70  	}
    71  	return appendNewline(b)
    72  }
    73  
    74  type rawProgressFormatter struct{}
    75  
    76  func (sf *rawProgressFormatter) formatStatus(id, format string, a ...interface{}) []byte {
    77  	return []byte(fmt.Sprintf(format, a...) + streamNewline)
    78  }
    79  
    80  func (sf *rawProgressFormatter) formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte {
    81  	if progress == nil {
    82  		progress = &jsonmessage.JSONProgress{}
    83  	}
    84  	endl := "\r"
    85  	if progress.String() == "" {
    86  		endl += "\n"
    87  	}
    88  	return []byte(action + " " + progress.String() + endl)
    89  }
    90  
    91  // NewProgressOutput returns a progress.Output object that can be passed to
    92  // progress.NewProgressReader.
    93  func NewProgressOutput(out io.Writer) progress.Output {
    94  	return &progressOutput{sf: &rawProgressFormatter{}, out: out, newLines: true}
    95  }
    96  
    97  // NewJSONProgressOutput returns a progress.Output that formats output
    98  // using JSON objects
    99  func NewJSONProgressOutput(out io.Writer, newLines bool) progress.Output {
   100  	return &progressOutput{sf: &jsonProgressFormatter{}, out: out, newLines: newLines}
   101  }
   102  
   103  type formatProgress interface {
   104  	formatStatus(id, format string, a ...interface{}) []byte
   105  	formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte
   106  }
   107  
   108  type progressOutput struct {
   109  	sf       formatProgress
   110  	out      io.Writer
   111  	newLines bool
   112  }
   113  
   114  // WriteProgress formats progress information from a ProgressReader.
   115  func (out *progressOutput) WriteProgress(prog progress.Progress) error {
   116  	var formatted []byte
   117  	if prog.Message != "" {
   118  		formatted = out.sf.formatStatus(prog.ID, prog.Message)
   119  	} else {
   120  		jsonProgress := jsonmessage.JSONProgress{Current: prog.Current, Total: prog.Total, HideCounts: prog.HideCounts, Units: prog.Units}
   121  		formatted = out.sf.formatProgress(prog.ID, prog.Action, &jsonProgress, prog.Aux)
   122  	}
   123  	_, err := out.out.Write(formatted)
   124  	if err != nil {
   125  		return err
   126  	}
   127  
   128  	if out.newLines && prog.LastUpdate {
   129  		_, err = out.out.Write(out.sf.formatStatus("", ""))
   130  		return err
   131  	}
   132  
   133  	return nil
   134  }
   135  
   136  // AuxFormatter is a streamFormatter that writes aux progress messages
   137  type AuxFormatter struct {
   138  	io.Writer
   139  }
   140  
   141  // Emit emits the given interface as an aux progress message
   142  func (sf *AuxFormatter) Emit(id string, aux interface{}) error {
   143  	auxJSONBytes, err := json.Marshal(aux)
   144  	if err != nil {
   145  		return err
   146  	}
   147  	auxJSON := new(json.RawMessage)
   148  	*auxJSON = auxJSONBytes
   149  	msgJSON, err := json.Marshal(&jsonmessage.JSONMessage{ID: id, Aux: auxJSON})
   150  	if err != nil {
   151  		return err
   152  	}
   153  	msgJSON = appendNewline(msgJSON)
   154  	n, err := sf.Writer.Write(msgJSON)
   155  	if n != len(msgJSON) {
   156  		return io.ErrShortWrite
   157  	}
   158  	return err
   159  }