find_ngrams_on_ct revision fe17456d5e528078ce69b5f15cf7adf1fab963f9 
1#!/usr/bin/env python
2
3
4"""Run Cluster Telemetry to compute n-grams from SKPs."""
5
6
7import argparse
8import os
9import re
10import subprocess
11import sys
12import tempfile
13
14
15NGRAMS_LUA_SUBSTITUTION_STR = '-- CHANGEME\nlocal n = \d+\n-- CHANGEME'
16SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
17
18
19def 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()
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
60if __name__ == '__main__':
61  main()
62