github.com/haraldrudell/parl@v0.4.176/pio/line-filter-writer.go (about)

     1  /*
     2  © 2023–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
     3  ISC License
     4  */
     5  
     6  package pio
     7  
     8  import (
     9  	"io"
    10  	"io/fs"
    11  	"sync"
    12  
    13  	"github.com/haraldrudell/parl"
    14  	"github.com/haraldrudell/parl/perrors"
    15  	"golang.org/x/exp/slices"
    16  )
    17  
    18  const (
    19  	newLineWriter = byte('\n')
    20  )
    21  
    22  // LineFilterFunc receives lines as they are written to the writer
    23  //   - can modify the line
    24  //   - can skip the line using skipLine
    25  //   - can return error
    26  type LineFilterFunc func(line *[]byte, isLastLine bool) (skipLine bool, err error)
    27  
    28  // LineFilterWriter is a writer that filters each line using a filter function
    29  type LineFilterWriter struct {
    30  	writeCloser io.WriteCloser
    31  	filter      LineFilterFunc
    32  
    33  	dataLock sync.Mutex
    34  	isClosed bool
    35  	data     []byte
    36  }
    37  
    38  var _ io.WriteCloser = &LineFilterWriter{}
    39  
    40  // NewLineFilterWriter is a writer that filters each line using a filter function
    41  func NewLineFilterWriter(writeCloser io.WriteCloser, filter LineFilterFunc) (lineWriter *LineFilterWriter) {
    42  	if writeCloser == nil {
    43  		panic(parl.NilError("writeCloser"))
    44  	} else if filter == nil {
    45  		panic(parl.NilError("filter"))
    46  	}
    47  	return &LineFilterWriter{writeCloser: writeCloser, filter: filter}
    48  }
    49  
    50  // Write saves data in slice and returns all bytes written or ErrFileAlreadyClosed
    51  func (wc *LineFilterWriter) Write(p []byte) (n int, err error) {
    52  	wc.dataLock.Lock()
    53  	defer wc.dataLock.Unlock()
    54  
    55  	if wc.isClosed {
    56  		err = perrors.ErrorfPF("%w", fs.ErrClosed)
    57  		return // closed return
    58  	}
    59  
    60  	// consume data
    61  	length := len(p)
    62  	for n < length {
    63  		index := slices.Index(p[n:], newLineWriter)
    64  
    65  		// check for p ending without newline
    66  		if index == -1 {
    67  			wc.data = append(wc.data, p[n:]...) // save in buffer
    68  			n = length                          // pretend data was written
    69  			break
    70  		}
    71  
    72  		index += n + 1 // include newline, make index in p
    73  		wc.data = append(wc.data, p[n:index]...)
    74  		if err = wc.processLine(); err != nil {
    75  			return
    76  		}
    77  		wc.data = wc.data[:0]
    78  		n = index
    79  	}
    80  
    81  	return // good write return
    82  }
    83  
    84  func (w *LineFilterWriter) processLine() (err error) {
    85  
    86  	// apply filter
    87  	if w.filter != nil {
    88  		var skipLine bool
    89  		if skipLine, err = w.invokeFilter(); err != nil || skipLine {
    90  			return
    91  		}
    92  	}
    93  
    94  	// write line to writeCloser
    95  	length := len(w.data)
    96  	var n int
    97  	for n < length {
    98  		var n0 int
    99  		if n0, err = w.writeCloser.Write(w.data[n:]); err != nil {
   100  			return
   101  		}
   102  		n += n0
   103  	}
   104  
   105  	return
   106  }
   107  
   108  // Close closes
   109  func (wc *LineFilterWriter) Close() (err error) {
   110  	wc.dataLock.Lock()
   111  	defer wc.dataLock.Unlock()
   112  
   113  	if wc.isClosed {
   114  		err = perrors.ErrorfPF("%w", fs.ErrClosed)
   115  		return // closed return
   116  	}
   117  
   118  	wc.isClosed = true
   119  	if len(wc.data) > 0 {
   120  		err = wc.processLine()
   121  	}
   122  
   123  	return
   124  }
   125  
   126  // invokeFilter captures a panic in the filter function
   127  func (w *LineFilterWriter) invokeFilter() (skipLine bool, err error) {
   128  	defer parl.RecoverErr(func() parl.DA { return parl.A() }, &err)
   129  
   130  	skipLine, err = w.filter(&w.data, w.isClosed)
   131  
   132  	return
   133  }