github.com/fluhus/gostuff@v0.4.1-0.20240331134726-be71864f2b5d/aio/open.go (about)

     1  // Package aio provides buffered file I/O.
     2  package aio
     3  
     4  import (
     5  	"bufio"
     6  	"compress/gzip"
     7  	"io"
     8  	"os"
     9  	"path/filepath"
    10  )
    11  
    12  const (
    13  	// If true, .gz files are automatically compressed/decompressed.
    14  	gzipSupport = true
    15  )
    16  
    17  // OpenRaw opens a file for reading, with a buffer.
    18  func OpenRaw(file string) (*Reader, error) {
    19  	f, err := os.Open(file)
    20  	if err != nil {
    21  		return nil, err
    22  	}
    23  	return &Reader{*bufio.NewReader(f), f}, nil
    24  }
    25  
    26  // CreateRaw opens a file for writing, with a buffer.
    27  // Erases any previously existing content.
    28  func CreateRaw(file string) (*Writer, error) {
    29  	f, err := os.Create(file)
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  	return &Writer{*bufio.NewWriter(f), f}, nil
    34  }
    35  
    36  // AppendRaw opens a file for writing, with a buffer.
    37  // Appends to previously existing content if any.
    38  func AppendRaw(file string) (*Writer, error) {
    39  	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0o644)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	return &Writer{*bufio.NewWriter(f), f}, nil
    44  }
    45  
    46  var (
    47  	rsuffixes = map[string]func(io.Reader) (io.Reader, error){}
    48  	wsuffixes = map[string]func(io.WriteCloser) (io.WriteCloser, error){}
    49  )
    50  
    51  // Open opens a file for reading, with a buffer.
    52  // Decompresses the data according to the file's suffix.
    53  func Open(file string) (*Reader, error) {
    54  	f, err := OpenRaw(file)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  	fn := rsuffixes[filepath.Ext(file)]
    59  	if fn == nil {
    60  		return f, nil
    61  	}
    62  	ff, err := fn(f)
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  	return &Reader{*bufio.NewReader(ff), f}, nil
    67  }
    68  
    69  // Create opens a file for writing, with a buffer.
    70  // Erases any previously existing content.
    71  // Compresses the data according to the file's suffix.
    72  func Create(file string) (*Writer, error) {
    73  	f, err := CreateRaw(file)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	fn := wsuffixes[filepath.Ext(file)]
    78  	if fn == nil {
    79  		return f, nil
    80  	}
    81  	ff, err := fn(f)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  	wrapper := &writerWrapper{ff, f}
    86  	return &Writer{*bufio.NewWriter(ff), wrapper}, nil
    87  }
    88  
    89  // Append opens a file for writing, with a buffer.
    90  // Appends to previously existing content if any.
    91  // Compresses the data according to the file's suffix.
    92  func Append(file string) (*Writer, error) {
    93  	f, err := AppendRaw(file)
    94  	if err != nil {
    95  		return nil, err
    96  	}
    97  	fn := wsuffixes[filepath.Ext(file)]
    98  	if fn == nil {
    99  		return f, nil
   100  	}
   101  	ff, err := fn(f)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  	wrapper := &writerWrapper{ff, f}
   106  	return &Writer{*bufio.NewWriter(ff), wrapper}, nil
   107  }
   108  
   109  // AddReadSuffix adds a supported suffix for automatic decompression.
   110  // suffix should include the dot. f should take a raw reader and return a reader
   111  // that decompresses the data.
   112  func AddReadSuffix(suffix string, f func(io.Reader) (io.Reader, error)) {
   113  	rsuffixes[suffix] = f
   114  }
   115  
   116  // AddWriteSuffix adds a supported suffix for automatic compression.
   117  // suffix should include the dot. f should take a raw writer and return a writer
   118  // that compresses the data.
   119  func AddWriteSuffix(suffix string, f func(io.WriteCloser) (
   120  	io.WriteCloser, error)) {
   121  	wsuffixes[suffix] = f
   122  }
   123  
   124  func init() {
   125  	if gzipSupport {
   126  		AddReadSuffix(".gz", func(r io.Reader) (io.Reader, error) {
   127  			z, err := gzip.NewReader(r)
   128  			return z, err
   129  		})
   130  		AddWriteSuffix(".gz", func(w io.WriteCloser) (io.WriteCloser, error) {
   131  			z, err := gzip.NewWriterLevel(w, 1)
   132  			return z, err
   133  		})
   134  	}
   135  }