github.com/alkemics/goflow@v0.2.1/gfutil/gfgo/writer.go (about)

     1  package gfgo
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  
     8  	"golang.org/x/tools/imports"
     9  
    10  	"mvdan.cc/gofumpt/format"
    11  )
    12  
    13  // Writer wraps another io.Writer, adding some Go code formatting through
    14  // golang.org/x/tools/imports.
    15  //
    16  // Because most of the time when writing to an io.Writer those writes happen
    17  // by batch (i.e. not at once), this Writer holds an internal bytes.Buffer
    18  // that the Write method uses. Use the Flush method to format the data and
    19  // actually write it to the underlying io.Writer.
    20  type Writer struct {
    21  	filename string
    22  	buf      bytes.Buffer
    23  	w        io.Writer
    24  }
    25  
    26  // NewWriter creates a new Writer that wraps w.
    27  //
    28  // The filename argument is passed to golang.org/x/tools/imports.Process when
    29  // flushing.
    30  func NewWriter(w io.Writer, filename string) *Writer {
    31  	return &Writer{
    32  		filename: filename,
    33  		buf:      bytes.Buffer{},
    34  		w:        w,
    35  	}
    36  }
    37  
    38  // Write writes into the bytes.Buffer held by the Writer.
    39  func (w *Writer) Write(src []byte) (int, error) {
    40  	return w.buf.Write(src)
    41  }
    42  
    43  // Flush formats the buffered data with golang.org/x/tools/imports.Process,
    44  // then writes it to the wrapped io.Writer.
    45  func (w *Writer) Flush() error {
    46  	src, err := imports.Process("", w.buf.Bytes(), nil)
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	src, err = format.Source(src, format.Options{LangVersion: "1.17"})
    52  	if err != nil {
    53  		return err
    54  	}
    55  
    56  	_, err = w.w.Write(src)
    57  	if err == nil {
    58  		w.buf.Reset()
    59  	}
    60  	return err
    61  }
    62  
    63  // The DebugWriter wraps an io.Writer, printing via fmt.Println
    64  // in case there is an error writing to the wrapped io.Writer.
    65  type DebugWriter struct {
    66  	Writer io.Writer
    67  }
    68  
    69  func (w DebugWriter) Write(src []byte) (int, error) {
    70  	l, err := w.Writer.Write(src)
    71  	if err != nil {
    72  		fmt.Println(string(src))
    73  	}
    74  	return l, err
    75  }