OLD | NEW |
---|---|
(Empty) | |
1 # Copyright 2016 The LUCI Authors. All rights reserved. | |
2 # Use of this source code is governed under the Apache License, Version 2.0 | |
3 # that can be found in the LICENSE file. | |
4 | |
5 """Command line tool for creating and extracting ar files.""" | |
6 | |
7 from __future__ import print_function | |
8 | |
9 import argparse | |
10 import io | |
11 import os | |
12 import shutil | |
13 import stat | |
14 import sys | |
15 import time | |
16 | |
17 # pylint: disable=relative-import | |
18 import arfile | |
19 | |
20 | |
21 class ProgressReporter(object): | |
22 def __init__(self, every): | |
23 self.every = int(every) | |
24 self.start = time.time() | |
25 self.filecount = 0 | |
26 self.lastreport = 0 | |
27 | |
28 def inc(self): | |
29 self.filecount += 1 | |
30 if (self.filecount - self.lastreport) >= self.every: | |
31 self.report() | |
32 | |
33 def report(self): | |
34 if self.every: | |
35 t = time.time()-self.start | |
36 print(u'Took %f for %i files == %f files/second' % ( | |
37 t, self.filecount, self.filecount/t), file=sys.stderr) | |
38 self.lastreport = self.filecount | |
39 | |
40 def __del__(self): | |
41 self.report() | |
42 | |
43 | |
44 def create_cmd(filename, dirs, progress, read_ahead, verbose): | |
45 afw = arfile.ArFileWriter(filename) | |
46 try: | |
47 for path in dirs: | |
48 for dirpath, child_dirs, filenames in os.walk(path): | |
49 # In-place sort the child_dirs so we walk in lexicographical order | |
50 child_dirs.sort() | |
51 filenames.sort() | |
52 for fn in filenames: | |
53 fp = os.path.join(dirpath, fn) | |
54 | |
55 if verbose: | |
56 print(fp, file=sys.stderr) | |
57 | |
58 with open(fp, 'rb') as f: | |
59 # If a file is small, it is cheaper to just read the file rather | |
60 # than doing a stat | |
61 data = f.read(read_ahead) | |
62 if len(data) < read_ahead: | |
63 afw.addfile(arfile.ArInfo.fromdefault( | |
64 fp[len(path)+1:], len(data)), io.BytesIO(data)) | |
65 else: | |
66 size = os.stat(fp).st_size | |
67 f.seek(0) | |
68 afw.addfile(arfile.ArInfo.fromdefault( | |
69 fp[len(path)+1:], size), f) | |
70 | |
71 progress.inc() | |
72 finally: | |
73 afw.close() | |
74 | |
75 | |
76 def list_cmd(filename, progress): | |
77 afr = arfile.ArFileReader(filename, fullparse=False) | |
78 for ai, _ in afr: | |
79 print(ai.name) | |
80 progress.inc() | |
81 | |
82 | |
83 def extract_cmd( | |
84 filename, progress, verbose, blocksize=1024*64): | |
85 afr = arfile.ArFileReader(filename, fullparse=False) | |
86 for ai, ifd in afr: | |
87 assert not ai.name.startswith('/') | |
88 if verbose: | |
89 print(ai.name, file=sys.stderr) | |
90 | |
91 try: | |
92 os.makedirs(os.path.dirname(ai.name)) | |
93 except OSError: | |
94 pass | |
95 | |
96 with open(ai.name, 'wb') as ofd: | |
97 written = 0 | |
98 while written < ai.size: | |
99 readsize = min(blocksize, ai.size-written) | |
100 ofd.write(ifd.read(readsize)) | |
101 written += readsize | |
102 | |
103 progress.inc() | |
104 | |
105 | |
106 def main(name, args): | |
107 parser = argparse.ArgumentParser( | |
108 prog=name, | |
109 description=sys.modules[__name__].__doc__) | |
110 subparsers = parser.add_subparsers( | |
111 dest='mode', help='sub-command help') | |
112 | |
113 # Create command | |
114 parser_create = subparsers.add_parser( | |
115 'create', help='Create a new ar file') | |
116 parser_create.add_argument( | |
117 "-r", "--read-ahead", | |
118 type=int, default=1024*64, | |
119 help="Amount of data to read-ahead before doing a stat.") | |
M-A Ruel
2016/06/22 14:21:32
use single quotes thoroughly.
mithro
2016/06/23 07:10:54
Done.
| |
120 parser_create.add_argument( | |
121 "-f", "--filename", | |
122 type=argparse.FileType('wb'), default=sys.stdout, | |
123 help="ar file to use") | |
124 parser_create.add_argument( | |
125 'dirs', nargs='+', help='Directory or file to add to the ar file') | |
126 | |
127 # List command | |
128 parser_list = subparsers.add_parser('list', help='List a new ar file') | |
129 | |
130 # Extract command | |
131 parser_extract = subparsers.add_parser( | |
132 'extract', help='Extract an existing ar file to current directory') | |
133 | |
134 # Add to output commands | |
135 for p in parser_list, parser_extract: | |
136 p.add_argument( | |
137 "-f", "--filename", | |
138 type=argparse.FileType('rb'), default=sys.stdin, | |
139 help="ar file to use") | |
140 | |
141 for p in parser_create, parser_extract: | |
142 p.add_argument( | |
143 "-v", "--verbose", | |
144 action="store_true", | |
145 help='Output file names to stderr while running.') | |
146 | |
147 # Add to all commands | |
148 for p in parser_create, parser_list, parser_extract: | |
149 p.add_argument( | |
150 "-p", "--progress", | |
151 type=ProgressReporter, default='10000', | |
152 help='Output progress information every N files.') | |
153 | |
154 args = parser.parse_args(args) | |
155 mode = getattr(sys.modules[__name__], args.mode + '_cmd') | |
156 del args.mode | |
157 return mode(**args.__dict__) | |
158 | |
159 | |
160 if __name__ == "__main__": | |
161 sys.exit(main("artool", (a.decode('utf-8') for a in sys.argv[1:]))) | |
OLD | NEW |