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