1"""util.py - General utilities for running, loading, and processing benchmarks 2""" 3import json 4import os 5import tempfile 6import subprocess 7import sys 8 9# Input file type enumeration 10IT_Invalid = 0 11IT_JSON = 1 12IT_Executable = 2 13 14_num_magic_bytes = 2 if sys.platform.startswith('win') else 4 15def 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 40def 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 54def 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 75def 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 88def 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 97def 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 117def 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