1 files changed, 61 insertions, 0 deletions
diff --git a/tools/lua/find_ngrams_on_ct b/tools/lua/find_ngrams_on_ct
new file mode 100755
index 0000000000..be3e4c07f1
--- /dev/null
+++ b/tools/lua/find_ngrams_on_ct
@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+
+
+"""Run Cluster Telemetry to compute n-grams from SKPs."""
+
+
+import argparse
+import os
+import re
+import subprocess
+import sys
+import tempfile
+
+
+NGRAMS_LUA_SUBSTITUTION_STR = '-- CHANGEME\nlocal n = \d+\n-- CHANGEME'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+
+
+def main():
+  # Parse arguments.
+  parser = argparse.ArgumentParser(
+      description='Run Cluster Telemetry to compute n-grams from SKPs.')
+  parser.add_argument('--n', help='Compute n-grams with this integer as N',
+                      required=True)
+  args, extra_args = parser.parse_known_args()
+
+  # Read the n-gram Lua script.
+  script_path = os.path.join(SCRIPT_DIR, 'ngrams.lua')
+  with open(script_path) as f:
+    script_contents = f.read()
+
+  # Replace the default value of n with the value specified by the user.
+  new_contents, subd = re.subn(NGRAMS_LUA_SUBSTITUTION_STR,
+                               'local n = %s' % args.n, script_contents, 1)
+  if subd != 1:
+    raise Exception('Unable to replace N in %s; expected to find:\n%s' % (
+                    script_path, sub))
+
+  # Write the new script contents to a temporary file.
+  tmp_script = tempfile.NamedTemporaryFile(delete=False)
+  try:
+    tmp_script.write(new_contents)
+    tmp_script.close()
+
+    # Run trigger_ct_lua with the new script, forwarding the rest of the
+    # passed-in arguments to the script.
+    trigger_ct_lua = os.path.join(SCRIPT_DIR, 'trigger_ct_lua')
+    cmd = ['python', trigger_ct_lua,
+           '--script', tmp_script.name,
+           '--aggregator', os.path.join(SCRIPT_DIR, 'ngrams_aggregate.lua')]
+    cmd.extend(extra_args)
+    try:
+      subprocess.check_call(cmd)
+    except subprocess.CalledProcessError:
+      exit(1)
+  finally:
+    os.remove(tmp_script.name)
+
+
+if __name__ == '__main__':
+  main()