OLD | NEW |
---|---|
(Empty) | |
1 #!/usr/bin/env python | |
2 | |
3 | |
4 """Run Cluster Telemetry to compute n-grams from SKPs.""" | |
5 | |
6 | |
7 import argparse | |
8 import os | |
9 import re | |
10 import subprocess | |
11 import sys | |
12 import tempfile | |
13 | |
14 | |
15 NGRAMS_LUA_SUBSTITUTION_STR = '-- CHANGEME\nlocal n = \d+\n-- CHANGEME' | |
16 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) | |
17 | |
18 | |
19 def main(): | |
20 # Parse arguments. | |
21 parser = argparse.ArgumentParser( | |
22 description='Run Cluster Telemetry to compute n-grams from SKPs.') | |
23 parser.add_argument('--n', help='Compute n-grams with this integer as N', | |
24 required=True) | |
25 args, extra_args = parser.parse_known_args() | |
26 | |
27 # Read the n-gram Lua script. | |
28 script_path = os.path.join(SCRIPT_DIR, 'ngrams.lua') | |
29 with open(script_path) as f: | |
30 script_contents = f.read() | |
31 | |
32 # Replace the default value of n with the value specified by the user. | |
33 new_contents, subd = re.subn(NGRAMS_LUA_SUBSTITUTION_STR, | |
34 'local n = %s' % args.n, script_contents, 1) | |
35 if subd != 1: | |
36 raise Exception('Unable to replace N in %s; expected to find:\n%s' % ( | |
37 script_path, sub)) | |
38 | |
39 # Write the new script contents to a temporary file. | |
40 tmp_script = tempfile.NamedTemporaryFile(delete=False) | |
41 try: | |
42 tmp_script.write(new_contents) | |
43 tmp_script.close() | |
rmistry
2015/06/30 17:05:12
Move the close into the finally?
borenet
2015/06/30 17:16:02
Doesn't matter for Unix, but on Windows we won't b
| |
44 | |
45 # Run trigger_ct_lua with the new script, forwarding the rest of the | |
46 # passed-in arguments to the script. | |
47 trigger_ct_lua = os.path.join(SCRIPT_DIR, 'trigger_ct_lua') | |
48 cmd = ['python', trigger_ct_lua, | |
49 '--script', tmp_script.name, | |
50 '--aggregator', os.path.join(SCRIPT_DIR, 'ngrams_aggregate.lua')] | |
51 cmd.extend(extra_args) | |
52 try: | |
53 subprocess.check_call(cmd) | |
54 except subprocess.CalledProcessError: | |
55 exit(1) | |
56 finally: | |
57 os.remove(tmp_script.name) | |
58 | |
59 | |
60 if __name__ == '__main__': | |
61 main() | |
OLD | NEW |