File: compress.go

package info (click to toggle)
golang-github-pierrec-lz4 4.1.18-1
  • links: PTS, VCS
  • area: main
  • in suites: experimental, sid, trixie
  • size: 89,368 kB
  • sloc: asm: 791; makefile: 6
file content (123 lines) | stat: -rw-r--r-- 2,928 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"sync/atomic"

	"code.cloudfoundry.org/bytefmt"
	"github.com/schollz/progressbar/v3"

	"github.com/pierrec/cmdflag"
	"github.com/pierrec/lz4/v4"
)

// Compress compresses a set of files or from stdin to stdout.
func Compress(fs *flag.FlagSet) cmdflag.Handler {
	var blockMaxSize string
	fs.StringVar(&blockMaxSize, "size", "4M", "block max size [64K,256K,1M,4M]")
	var blockChecksum bool
	fs.BoolVar(&blockChecksum, "bc", false, "enable block checksum")
	var streamChecksum bool
	fs.BoolVar(&streamChecksum, "sc", false, "disable stream checksum")
	var level int
	fs.IntVar(&level, "l", 0, "compression level (0=fastest)")
	var concurrency int
	fs.IntVar(&concurrency, "c", -1, "concurrency (default=all CPUs")

	return func(args ...string) (int, error) {
		sz, err := bytefmt.ToBytes(blockMaxSize)
		if err != nil {
			return 0, err
		}

		zw := lz4.NewWriter(nil)
		options := []lz4.Option{
			lz4.BlockChecksumOption(blockChecksum),
			lz4.BlockSizeOption(lz4.BlockSize(sz)),
			lz4.ChecksumOption(streamChecksum),
			lz4.CompressionLevelOption(lz4.CompressionLevel(level)),
			lz4.ConcurrencyOption(concurrency),
		}
		if err := zw.Apply(options...); err != nil {
			return 0, err
		}

		// Use stdin/stdout if no file provided.
		if len(args) == 0 {
			zw.Reset(os.Stdout)
			_, err := io.Copy(zw, os.Stdin)
			if err != nil {
				return 0, err
			}
			return 0, zw.Close()
		}

		for fidx, filename := range args {
			// Input file.
			file, err := os.Open(filename)
			if err != nil {
				return fidx, err
			}
			finfo, err := file.Stat()
			if err != nil {
				return fidx, err
			}
			mode := finfo.Mode() // use the same mode for the output file

			// Accumulate compressed bytes num.
			var (
				zsize int64
				size  = finfo.Size()
			)
			if size > 0 {
				// Progress bar setup.
				numBlocks := int(size) / int(sz)
				bar := progressbar.NewOptions(numBlocks,
					// File transfers are usually slow, make sure we display the bar at 0%.
					progressbar.OptionSetRenderBlankState(true),
					// Display the filename.
					progressbar.OptionSetDescription(filename),
					progressbar.OptionClearOnFinish(),
				)
				err = zw.Apply(
					lz4.OnBlockDoneOption(func(n int) {
						_ = bar.Add(1)
						atomic.AddInt64(&zsize, int64(n))
					}),
				)
				if err != nil {
					return 0, err
				}
			}

			// Output file.
			zfilename := fmt.Sprintf("%s%s", filename, lz4Extension)
			zfile, err := os.OpenFile(zfilename, os.O_CREATE|os.O_WRONLY, mode)
			if err != nil {
				return fidx, err
			}
			zw.Reset(zfile)

			// Compress.
			_, err = io.Copy(zw, file)
			if err != nil {
				return fidx, err
			}
			for _, c := range []io.Closer{zw, zfile} {
				err := c.Close()
				if err != nil {
					return fidx, err
				}
			}

			if size > 0 {
				fmt.Printf("%s %.02f%%\n", zfilename, float64(zsize)*100/float64(size))
			}
		}

		return len(args), nil
	}
}