aboutsummaryrefslogtreecommitdiff
path: root/scripts/zstd.py
blob: f85e7dae3cb6e6a77a9fdff9aed4f8e3daee0ea1 (plain)
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
import subprocess
from rcc_format import *

def zstd(filelist, args):
    if args.train is not None:
        train(filelist, args.train, args.binary, args.dsize)
        return

    write_header(args.output, args.namespace)

    for f in filelist:
        with open(f.path, 'rb') as contents:
            write_item(args.output, f.variable, compress(contents, args.binary, args.level, dictionary=args.dict))
    
    write_entries(args.output, filelist)
    if args.dict is not None:
        write_item(args.output, 'dict', args.dict.read())
        print("constexpr auto dictionary = std::span(dict, dict_len);", file=args.output)
    
    print("constexpr auto compression = embed::Zstd;", file=args.output)
    
    write_footer(args.output, args.namespace)

def train(filelist, output, zstd_bin, maxdict):
    cmd = [ zstd_bin, '--train', '--maxdict=' + str(maxdict), '-o', output.name ]

    for f in filelist:
        cmd.append(f.path)

    subprocess.run(cmd)

def compress(file, zstd_bin, level, dictionary=None):
    cmd = [ zstd_bin, '--compress', '--stdout', '-' + str(level) ]

    if dictionary is not None:
        cmd.append('-D')
        cmd.append(dictionary.name)
    
    cmd.append(file.name)
    return subprocess.run(cmd, capture_output=True).stdout