Improve the profiling script

scripts/profile_output.py

@@ -1,10 +1,10 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3.6
 """
-Profile a piece of Python code with ``cProfile``. Tries a completion on a
+Profile a piece of Python code with ``profile``. Tries a completion on a
 certain piece of code.
 
 Usage:
-  profile.py [<code>] [-n <number>] [-d] [-o] [-s <sort>]
+  profile.py [<code>] [-n <number>] [-d] [-o] [-s <sort>] [-i]
   profile.py -h | --help
 
 Options:
@@ -12,34 +12,41 @@ Options:
   -n <number>   Number of passes before profiling [default: 1].
   -d --debug    Enable Jedi internal debugging.
   -o --omit     Omit profiler, just do a normal run.
+  -i --infer    Infer types instead of completions.
   -s <sort>     Sort the profile results, e.g. cum, name [default: time].
 """
 
 import time
-import cProfile
+import profile
 
 from docopt import docopt
 import jedi
 
 
-def run(code, index):
+def run(code, index, infer=False):
     start = time.time()
-    result = jedi.Script(code).completions()
+    script = jedi.Script(code)
+    if infer:
+        result = script.goto_definitions()
+    else:
+        result = script.completions()
     print('Used %ss for the %sth run.' % (time.time() - start, index + 1))
     return result
 
 
 def main(args):
     code = args['<code>']
+    infer = args['--infer']
     n = int(args['-n'])
+
     for i in range(n):
-        run(code, i)
+        run(code, i, infer=infer)
 
     jedi.set_debug_function(notices=args['--debug'])
     if args['--omit']:
-        run(code, n)
+        run(code, n, infer=infer)
     else:
-        cProfile.runctx('run(code, n)', globals(), locals(), sort=args['-s'])
+        profile.runctx('run(code, n, infer=infer)', globals(), locals(), sort=args['-s'])
 
 
 if __name__ == '__main__':