github.com/grailbio/base@v0.0.11/recordio/recordioutil/flags.go (about)

     1  // Copyright 2017 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package recordioutil
     6  
     7  import (
     8  	"flag"
     9  	"fmt"
    10  
    11  	"github.com/klauspost/compress/flate"
    12  )
    13  
    14  // CompressionLevelFlag represents a command line flag to specify a compression
    15  // level
    16  type CompressionLevelFlag struct {
    17  	// Specified will be true if the flag has been specified on the command line.
    18  	Specified bool
    19  	// Level corresponds to the consts in the flate package
    20  	Level int
    21  }
    22  
    23  // Set implements flag.Value.
    24  func (f *CompressionLevelFlag) Set(v string) error {
    25  	switch v {
    26  	case "none":
    27  		f.Level = flate.NoCompression
    28  	case "fastest":
    29  		f.Level = flate.BestSpeed
    30  	case "best":
    31  		f.Level = flate.BestCompression
    32  	case "default":
    33  		f.Level = flate.DefaultCompression
    34  	case "huffman-only":
    35  		f.Level = flate.HuffmanOnly
    36  	default:
    37  		return fmt.Errorf("unrecognised compression option: %v: not one of none, fastest, best, default or huffman-only", v)
    38  	}
    39  	f.Specified = true
    40  	return nil
    41  }
    42  
    43  // String implements flag.Value.
    44  func (f *CompressionLevelFlag) String() string {
    45  	switch f.Level {
    46  	case flate.NoCompression:
    47  		return "none"
    48  	case flate.BestSpeed:
    49  		return "fastest"
    50  	case flate.BestCompression:
    51  		return "best"
    52  	case flate.DefaultCompression:
    53  		return "default"
    54  	case flate.HuffmanOnly:
    55  		return "huffman-only"
    56  	}
    57  	panic(fmt.Sprintf("unrecognised compression constant: %v", f.Level))
    58  }
    59  
    60  // WriterFlags represents the flags required to configure a recordioutil.Writer.
    61  type WriterFlags struct {
    62  	CompressionFlag    CompressionLevelFlag
    63  	MegaBytesPerRecord uint
    64  	ItemsPerRecord     uint
    65  }
    66  
    67  // WriterOptsFromFlags determines the WriterOpts to use based on the supplied
    68  // command line flags.
    69  func WriterOptsFromFlags(flags *WriterFlags) WriterOpts {
    70  	return WriterOpts{
    71  		MaxItems:   uint32(flags.ItemsPerRecord),
    72  		MaxBytes:   uint32(flags.MegaBytesPerRecord * 1024 * 1024),
    73  		FlateLevel: flags.CompressionFlag.Level,
    74  	}
    75  }
    76  
    77  // RegisterWriterFlags registers the flags in WriterFlags with the supplied
    78  // FlagSet. If the values of ItemsPerRecord, MegaBytesPerRecord or
    79  // CompressionFlag.Level have their go default values then appropriate defaults
    80  // will be used instead.
    81  func RegisterWriterFlags(flag *flag.FlagSet, opts *WriterFlags) {
    82  	if opts.ItemsPerRecord == 0 {
    83  		opts.ItemsPerRecord = 16
    84  	}
    85  	if opts.MegaBytesPerRecord == 0 {
    86  		opts.MegaBytesPerRecord = 1024 * 1024
    87  	}
    88  	if opts.CompressionFlag.Level == 0 {
    89  		opts.CompressionFlag.Level = flate.BestCompression
    90  	}
    91  	flag.Var(&opts.CompressionFlag, "recordio-compression-level", "compression level, one of: none|fastest|best|default|huffman-only")
    92  	flag.UintVar(&opts.MegaBytesPerRecord, "recordio-MiB-per-record", opts.MegaBytesPerRecord, "megabytes per recordio record")
    93  	flag.UintVar(&opts.ItemsPerRecord, "recordio-items-per-record", opts.ItemsPerRecord, "items per recordio record")
    94  }