github.com/mithrandie/csvq@v1.18.1/lib/file/reader.go (about)

     1  package file
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"os"
     7  	"strings"
     8  )
     9  
    10  type Reader struct {
    11  	offset    int64
    12  	reader    io.Reader
    13  	headBytes *bytes.Reader
    14  	headLen   int64
    15  }
    16  
    17  func NewReader(r io.Reader, headLen int) (*Reader, error) {
    18  	fileHead := make([]byte, headLen)
    19  	n, err := r.Read(fileHead)
    20  	if err != nil && err != io.EOF {
    21  		return nil, err
    22  	}
    23  
    24  	headBytes := bytes.NewReader(fileHead[:n])
    25  
    26  	return &Reader{
    27  		offset:    0,
    28  		reader:    r,
    29  		headBytes: headBytes,
    30  		headLen:   int64(n),
    31  	}, nil
    32  }
    33  
    34  func (r *Reader) Read(p []byte) (n int, err error) {
    35  	if r.headLen <= r.offset {
    36  		n, err = r.reader.Read(p)
    37  		if err == nil || err == io.EOF {
    38  			r.offset = r.offset + int64(n)
    39  		}
    40  		return
    41  	}
    42  
    43  	readLen := len(p)
    44  	_, err = r.headBytes.Seek(r.offset, io.SeekStart)
    45  	if err != nil {
    46  		return
    47  	}
    48  
    49  	n, err = r.headBytes.Read(p)
    50  	if err != nil {
    51  		return
    52  	}
    53  
    54  	delta := readLen - n
    55  	if delta == 0 {
    56  		r.offset = r.offset + int64(n)
    57  		return
    58  	}
    59  
    60  	b := make([]byte, delta)
    61  	n2, err := r.reader.Read(b)
    62  	if err == nil || err == io.EOF {
    63  		for i := 0; i < n2; i++ {
    64  			p[n+i] = b[i]
    65  		}
    66  	}
    67  	n = n + n2
    68  	r.offset = r.offset + int64(n)
    69  	return
    70  }
    71  
    72  func (r *Reader) HeadBytes() (io.ReadSeeker, error) {
    73  	_, err := r.headBytes.Seek(0, io.SeekStart)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  
    78  	fileHead := make([]byte, r.headLen)
    79  	_, err = r.headBytes.Read(fileHead)
    80  	if err != nil && err != io.EOF {
    81  		return nil, err
    82  	}
    83  
    84  	return bytes.NewReader(fileHead), nil
    85  }
    86  
    87  func (r *Reader) Size() int64 {
    88  	switch s := r.reader.(type) {
    89  	case *os.File:
    90  		if fi, err := s.Stat(); err == nil {
    91  			return fi.Size()
    92  		}
    93  	case *strings.Reader:
    94  		return s.Size()
    95  	case *bytes.Reader:
    96  		return s.Size()
    97  	}
    98  	return 0
    99  }