forked from pierrec/lz4
-
Notifications
You must be signed in to change notification settings - Fork 0
/
compress.go
115 lines (102 loc) · 2.66 KB
/
compress.go
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
package main
import (
"flag"
"fmt"
"io"
"os"
"sync/atomic"
"code.cloudfoundry.org/bytefmt"
"github.com/schollz/progressbar"
"github.com/pierrec/lz4"
"github.com/pierrec/cmdflag"
)
// 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)
zw.Header = lz4.Header{
BlockChecksum: blockChecksum,
BlockMaxSize: int(sz),
NoChecksum: streamChecksum,
CompressionLevel: level,
}
zw.WithConcurrency(concurrency)
// 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) / zw.Header.BlockMaxSize
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(),
)
zw.OnBlockDone = func(n int) {
_ = bar.Add(1)
atomic.AddInt64(&zsize, int64(n))
}
}
// Output file.
zfilename := fmt.Sprintf("%s%s", filename, lz4.Extension)
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
}
}