Home | History | Annotate | Download | only in gbench 
      1 """util.py - General utilities for running, loading, and processing benchmarks
      2 """
      3 import json
      4 import os
      5 import tempfile
      6 import subprocess
      7 import sys
      8 
      9 # Input file type enumeration
     10 IT_Invalid    = 0
     11 IT_JSON       = 1
     12 IT_Executable = 2
     13 
     14 _num_magic_bytes = 2 if sys.platform.startswith('win') else 4
     15 def is_executable_file(filename):
     16     """
     17     Return 'True' if 'filename' names a valid file which is likely
     18     an executable. A file is considered an executable if it starts with the
     19     magic bytes for a EXE, Mach O, or ELF file.
     20     """
     21     if not os.path.isfile(filename):
     22         return False
     23     with open(filename, 'r') as f:
     24         magic_bytes = f.read(_num_magic_bytes)
     25     if sys.platform == 'darwin':
     26         return magic_bytes in [
     27             '\xfe\xed\xfa\xce',  # MH_MAGIC
     28             '\xce\xfa\xed\xfe',  # MH_CIGAM
     29             '\xfe\xed\xfa\xcf',  # MH_MAGIC_64
     30             '\xcf\xfa\xed\xfe',  # MH_CIGAM_64
     31             '\xca\xfe\xba\xbe',  # FAT_MAGIC
     32             '\xbe\xba\xfe\xca'   # FAT_CIGAM
     33         ]
     34     elif sys.platform.startswith('win'):
     35         return magic_bytes == 'MZ'
     36     else:
     37         return magic_bytes == '\x7FELF'
     38 
     39 
     40 def is_json_file(filename):
     41     """
     42     Returns 'True' if 'filename' names a valid JSON output file.
     43     'False' otherwise.
     44     """
     45     try:
     46         with open(filename, 'r') as f:
     47             json.load(f)
     48         return True
     49     except:
     50         pass
     51     return False
     52 
     53 
     54 def classify_input_file(filename):
     55     """
     56     Return a tuple (type, msg) where 'type' specifies the classified type
     57     of 'filename'. If 'type' is 'IT_Invalid' then 'msg' is a human readable
     58     string represeting the error.
     59     """
     60     ftype = IT_Invalid
     61     err_msg = None
     62     if not os.path.exists(filename):
     63         err_msg = "'%s' does not exist" % filename
     64     elif not os.path.isfile(filename):
     65         err_msg = "'%s' does not name a file" % filename
     66     elif is_executable_file(filename):
     67         ftype = IT_Executable
     68     elif is_json_file(filename):
     69         ftype = IT_JSON
     70     else:
     71         err_msg = "'%s' does not name a valid benchmark executable or JSON file"
     72     return ftype, err_msg
     73 
     74 
     75 def check_input_file(filename):
     76     """
     77     Classify the file named by 'filename' and return the classification.
     78     If the file is classified as 'IT_Invalid' print an error message and exit
     79     the program.
     80     """
     81     ftype, msg = classify_input_file(filename)
     82     if ftype == IT_Invalid:
     83         print "Invalid input file: %s" % msg
     84         sys.exit(1)
     85     return ftype
     86 
     87 
     88 def load_benchmark_results(fname):
     89     """
     90     Read benchmark output from a file and return the JSON object.
     91     REQUIRES: 'fname' names a file containing JSON benchmark output.
     92     """
     93     with open(fname, 'r') as f:
     94         return json.load(f)
     95 
     96 
     97 def run_benchmark(exe_name, benchmark_flags):
     98     """
     99     Run a benchmark specified by 'exe_name' with the specified
    100     'benchmark_flags'. The benchmark is run directly as a subprocess to preserve
    101     real time console output.
    102     RETURNS: A JSON object representing the benchmark output
    103     """
    104     thandle, tname = tempfile.mkstemp()
    105     os.close(thandle)
    106     cmd = [exe_name] + benchmark_flags
    107     print("RUNNING: %s" % ' '.join(cmd))
    108     exitCode = subprocess.call(cmd + ['--benchmark_out=%s' % tname])
    109     if exitCode != 0:
    110         print('TEST FAILED...')
    111         sys.exit(exitCode)
    112     json_res = load_benchmark_results(tname)
    113     os.unlink(tname)
    114     return json_res
    115 
    116 
    117 def run_or_load_benchmark(filename, benchmark_flags):
    118     """
    119     Get the results for a specified benchmark. If 'filename' specifies
    120     an executable benchmark then the results are generated by running the
    121     benchmark. Otherwise 'filename' must name a valid JSON output file,
    122     which is loaded and the result returned.
    123     """
    124     ftype = check_input_file(filename)
    125     if ftype == IT_JSON:
    126         return load_benchmark_results(filename)
    127     elif ftype == IT_Executable:
    128         return run_benchmark(filename, benchmark_flags)
    129     else:
    130         assert False # This branch is unreachable