Commit | Line | Data |
---|---|---|
3b8cae5a ERE |
1 | #!/usr/bin/env python3 |
2 | ||
3 | # Copyright (C) 2014 Intra2net AG | |
4 | # | |
5 | # This program is free software; you can redistribute it and/or modify | |
6 | # it under the terms of the GNU Lesser General Public License as published | |
7 | # by the Free Software Foundation; either version 3 of the License, or | |
8 | # (at your option) any later version. | |
9 | # | |
10 | # This program is distributed in the hope that it will be useful, | |
11 | # but WITHOUT ANY WARRANTY; without even the implied warranty of | |
12 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the | |
13 | # GNU Lesser General Public License for more details. | |
14 | # | |
15 | # You should have received a copy of the GNU General Public License | |
16 | # along with this program. If not, see | |
17 | # <http://www.gnu.org/licenses/lgpl-3.0.html> | |
18 | ||
19 | # Authors: Victor Ramirez de la Corte <victor@wadobo.com> | |
20 | # Eduardo Robles Elvira <edulix@wadobo.com> | |
21 | ||
22 | import argparse | |
23 | import os | |
24 | import zlib | |
25 | import fnmatch | |
26 | import cProfile | |
27 | import io | |
5368092d | 28 | import binascii |
3b8cae5a ERE |
29 | import pstats |
30 | ||
31 | ||
08bb4b4a | 32 | def main(args = None): |
3b8cae5a ERE |
33 | ''' |
34 | Main function, parses the command line arguments and launches the appropiate | |
35 | benchmark test. | |
36 | ''' | |
37 | parser = argparse.ArgumentParser(description='Profiling test options. ') | |
904d3019 ERE |
38 | parser.add_argument('-t', '--test', |
39 | choices=['delta-tarfile', 'tarfile', 'zlib', | |
40 | 'zlib-blocks'], | |
3b8cae5a ERE |
41 | help='Select option for testing. ') |
42 | parser.add_argument('-l', '--compression-level', type=int, default=9, | |
43 | choices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], | |
44 | help='Select level of compression for zlib. ') | |
45 | parser.add_argument('-p', '--path', | |
46 | help='path to the file/dir to compress.') | |
47 | parser.add_argument('-P', '--profile', action='store_true', | |
48 | help='Enable profile') | |
49 | parser.add_argument('-m', '--tarmode', default="w:gz", | |
5cd23642 | 50 | help='Specify tarfile mode [w:gz]') |
904d3019 ERE |
51 | parser.add_argument('-st', '--strategy', default='default', |
52 | choices=['default', 'filtered', 'huffman_only'], | |
53 | help='Select the deflate strategy in zlib related ' | |
54 | 'tests.') | |
3b8cae5a ERE |
55 | parser.add_argument( |
56 | '-s', '--sort', | |
57 | help='Sort output profile', | |
58 | choices=[ 'calls', 'cumulative', 'cumtime', 'file', 'path', 'module', | |
59 | 'ncalls', 'pcalls', 'line', 'name', 'nfl', 'stdname', 'time', | |
60 | 'tottime'], | |
61 | default="cumulative") | |
62 | ||
08bb4b4a ERE |
63 | if not args: |
64 | pargs = parser.parse_args() | |
65 | else: | |
66 | pargs = parser.parse_args(args) | |
67 | ||
3b8cae5a ERE |
68 | compression_level = pargs.compression_level |
69 | path = pargs.path | |
70 | ||
71 | # start profiling | |
72 | if pargs.profile: | |
73 | pr = cProfile.Profile() | |
74 | pr.enable() | |
75 | ||
76 | if pargs.test in ['delta-tarfile', 'tarfile']: | |
77 | test_tarfile(pargs.test, compression_level, path, pargs.tarmode) | |
904d3019 ERE |
78 | elif pargs.test == 'zlib': |
79 | test_zlib(compression_level, path, pargs.strategy) | |
80 | elif pargs.test == 'zlib-blocks': | |
81 | test_zlib_blocks(compression_level, path, pargs.strategy) | |
3b8cae5a ERE |
82 | else: |
83 | parser.print_help() | |
84 | ||
85 | # end profiling | |
86 | if pargs.profile: | |
87 | pr.disable() | |
88 | print_profile(pr, pargs.sort) | |
89 | ||
90 | def test_tarfile(res, compression_level=9, path='source_dir', tarmode="w:gz"): | |
91 | ''' | |
5368092d | 92 | Test that creates a tarfile called test_tarfile.tar.gz from a directory. |
3b8cae5a | 93 | ''' |
5368092d PG |
94 | encryptor = None |
95 | dstfile = "test_tarfile.tar.gz" | |
96 | ||
3b8cae5a | 97 | if res == 'delta-tarfile': |
3b8cae5a | 98 | import sys |
5368092d PG |
99 | # need to fix up the path before we can load a module explicitly |
100 | dtarpath = os.getenv ("DELTATAR") | |
101 | if dtarpath is not None: | |
102 | sys.path.append (dtarpath) | |
103 | ||
104 | try: | |
105 | from deltatar import tarfile as TarFile | |
106 | from deltatar import crypto | |
107 | except ImportError: | |
108 | raise Exception ("location of deltatar not found; " | |
109 | "set $DELTATAR to provide an explicit path") | |
110 | ||
111 | if tarmode.endswith (".pdtcrypt"): | |
112 | dstfile = dstfile + ".pdtcrypt" | |
113 | tarmode = tarmode.rstrip (".pdtcrypt") | |
114 | encryptor = crypto.Encrypt (password="test1234", | |
115 | nacl=binascii.unhexlify(b"0011223344556677" | |
116 | b"8899aabbccddeeff"), | |
117 | version=1, | |
118 | paramversion=1) | |
119 | ||
3b8cae5a ERE |
120 | elif res == 'tarfile': |
121 | import tarfile as TarFile | |
122 | ||
5368092d PG |
123 | tar = TarFile.open (dstfile, |
124 | mode=tarmode, | |
125 | concat='#' in tarmode, | |
126 | encryption=encryptor, | |
127 | compresslevel=compression_level) | |
128 | tar.add(path) | |
129 | ||
3b8cae5a | 130 | |
904d3019 ERE |
131 | def test_zlib(compression_level=9, path='source_dir', strategy='default', |
132 | BUFSIZE = 16 * 1024): | |
133 | ''' | |
134 | Test that compresses a directory with zlib | |
135 | ''' | |
136 | if strategy == 'default': | |
137 | strategy = zlib.Z_DEFAULT_STRATEGY | |
138 | elif strategy == 'filtered': | |
139 | strategy = zlib.Z_FILTERED | |
140 | else: | |
141 | strategy = zlib.Z_HUFFMAN_ONLY | |
142 | ||
143 | c = zlib.compressobj(compression_level, | |
144 | zlib.DEFLATED, | |
145 | -zlib.MAX_WBITS, | |
146 | zlib.DEF_MEM_LEVEL, | |
147 | strategy) | |
148 | ||
149 | def add(path): | |
150 | l = 0 | |
151 | l_orig = 0 | |
152 | if os.path.isdir(path): | |
153 | for f in os.listdir(path): | |
154 | l2, l_orig2 = add(os.path.join(path, f)) | |
155 | l += l2 | |
156 | l_orig += l_orig2 | |
157 | elif os.path.isfile(path): | |
158 | with open(path, 'rb') as fo: | |
159 | while True: | |
160 | buf = fo.read(BUFSIZE) | |
161 | l_orig += len(buf) | |
162 | if len(buf) == 0: | |
163 | return l, l_orig | |
164 | buf = c.compress(buf) | |
165 | if buf is not None: | |
166 | l += len(buf) | |
167 | return l, l_orig | |
168 | ||
169 | l, l_orig = add(path) | |
170 | l += len(c.flush()) | |
171 | print("total size: %d MiB compressed, %d MiB uncompressed" % ( | |
172 | (l / (1024 ** 2)), | |
173 | (l_orig / (1024 ** 2)))) | |
174 | ||
175 | ||
176 | def test_zlib_blocks(compression_level=9, path='source_dir', strategy='default', | |
177 | BUFSIZE = 16 * 1024): | |
178 | ''' | |
179 | Test that compresses a directory with zlib | |
180 | ''' | |
181 | if strategy == 'default': | |
182 | strategy = zlib.Z_DEFAULT_STRATEGY | |
183 | elif strategy == 'filtered': | |
184 | strategy = zlib.Z_FILTERED | |
185 | else: | |
186 | strategy = zlib.Z_HUFFMAN_ONLY | |
187 | ||
188 | def add(path): | |
189 | l = 0 | |
190 | l_orig = 0 | |
191 | if os.path.isdir(path): | |
192 | for f in os.listdir(path): | |
193 | l2, l_orig2 = add(os.path.join(path, f)) | |
194 | l += l2 | |
195 | l_orig += l_orig2 | |
196 | elif os.path.isfile(path): | |
197 | c = zlib.compressobj(compression_level, | |
198 | zlib.DEFLATED, | |
199 | -zlib.MAX_WBITS, | |
200 | zlib.DEF_MEM_LEVEL, | |
201 | strategy) | |
202 | with open(path, 'rb') as fo: | |
203 | while True: | |
204 | buf = fo.read(BUFSIZE) | |
205 | l_orig += len(buf) | |
206 | if len(buf) == 0: | |
207 | l += len(c.flush()) | |
208 | return l, l_orig | |
209 | buf = c.compress(buf) | |
210 | if buf is not None: | |
211 | l += len(buf) | |
212 | return l, l_orig | |
213 | ||
214 | l, l_orig = add(path) | |
215 | print("total size: %d MiB compressed, %d MiB uncompressed" % ( | |
216 | (l / (1024 ** 2)), | |
217 | (l_orig / (1024 ** 2)))) | |
218 | ||
3b8cae5a ERE |
219 | def print_profile(pr, sort): |
220 | s = io.StringIO() | |
221 | if not sort: | |
222 | sort = 'cumulative' | |
223 | ps = pstats.Stats(pr, stream=s).sort_stats(sort) | |
224 | ps.print_stats(20) | |
225 | print(s.getvalue()) | |
226 | ||
227 | ||
228 | if __name__ == '__main__': | |
229 | main() |