Compression wrappers for io.Reader and io.Writer

compressed_reader_writer.go

// Copyright (c) 2024 Neomantra BV
//
// Opinionated Reader/Writer wrappers

package nmio

import (
	"compress/gzip"
	"io"
	"os"
	"strings"

	"github.com/klauspost/compress/zstd"
)

//////////////////////////////////////////////////////////////////////////////

type CompressionType int

const (
	CT_UNKNOWN CompressionType = iota // Compression format is unknown
	CT_NONE                           // Uncompress
	CT_ZSTD                           // Compression format is ZSTD
	CT_GZ                             // Commpresion format is GZIP
)

func IsFileGZ(filename string, compType CompressionType) bool {
	return (compType == CT_GZ) ||
		((compType == CT_UNKNOWN) && strings.HasSuffix(filename, ".gz"))
}

func IsFileZSTD(filename string, compType CompressionType) bool {
	return (compType == CT_ZSTD) ||
		((compType == CT_UNKNOWN) && (strings.HasSuffix(filename, ".zst") || strings.HasSuffix(filename, ".zstd")))
}

//////////////////////////////////////////////////////////////////////////////

// Returns a io.Reader for the given filename, or os.Stdout if filename is "-".
// If isGZ is true or the filename ends in ".gz", the writer will gzip the output.
func MakeCompressedReader(filename string, compType CompressionType) (io.Reader, io.Closer, error) {
	var reader io.Reader
	var closer io.Closer

	if filename != "-" {
		if file, err := os.Open(filename); err == nil {
			reader, closer = file, file
		} else {
			return nil, nil, err
		}
	} else {
		reader, closer = os.Stdin, nil
	}

	var err error
	if IsFileGZ(filename, compType) {
		reader, err = gzip.NewReader(reader)
	} else if IsFileZSTD(filename, compType) {
		reader, err = zstd.NewReader(reader)
	}

	if err != nil {
		// clean up file
		if closer != nil {
			closer.Close()
		}
		return nil, nil, err
	}
	return reader, closer, nil
}

// Returns an io.Writer for the given filename, or os.Stdout if filename is "-".  Also returns a closing function to defer and any error.
// If isGZ is true or the filename ends in ".gz", the writer will gzip-compress the output.
// If isZstd is true or the filename ends in ".zst" or ".zstd", the writer will zstd-compress the output.
func MakeCompressedWriter(filename string, compType CompressionType) (io.Writer, func(), error) {
	var writer io.Writer
	var closer io.Closer
	fileCloser := func() {
		if closer != nil {
			closer.Close()
		}
	}
	if filename != "-" {
		if file, err := os.Create(filename); err == nil {
			writer, closer = file, file
		} else {
			return nil, nil, err
		}
	} else {
		writer, closer = os.Stdout, nil
	}

	if IsFileGZ(filename, compType) {
		gzipWriter := gzip.NewWriter(writer)
		gzipCloser := func() {
			gzipWriter.Close()
			fileCloser()
		}
		return gzipWriter, gzipCloser, nil
	} else if IsFileZSTD(filename, compType) {
		zstdWriter, err := zstd.NewWriter(writer)
		if err != nil {
			fileCloser()
			return nil, nil, err
		}
		zstdCloser := func() {
			zstdWriter.Close()
			fileCloser()
		}
		return zstdWriter, zstdCloser, nil
	} else {
		return writer, fileCloser, nil
	}
}