diff --git a/bin/ycsb b/bin/ycsb index e85ca43d70..8947cbc5cd 100755 --- a/bin/ycsb +++ b/bin/ycsb @@ -23,11 +23,13 @@ import os import shlex import sys import subprocess + try: + mod = __import__('argparse') import argparse except ImportError: print >> sys.stderr, '[ERROR] argparse not found. Try installing it via "pip".' - raise + exit(1) BASE_URL = "https://github.com/brianfrankcooper/YCSB/tree/master/" COMMANDS = { @@ -122,12 +124,41 @@ def usage(): return output.getvalue() -def check_output(cmd): - p = subprocess.Popen(cmd, stdout=subprocess.PIPE) - stdout, _ = p.communicate() - if p.returncode: - raise subprocess.CalledProcessError(p.returncode, cmd) - return stdout +# Python 2.6 doesn't have check_output. Add the method as it is in Python 2.7 +# Based on https://github.com/python/cpython/blob/2.7/Lib/subprocess.py#L545 +def check_output(*popenargs, **kwargs): + r"""Run command with arguments and return its output as a byte string. + + If the exit code was non-zero it raises a CalledProcessError. The + CalledProcessError object will have the return code in the returncode + attribute and output in the output attribute. + + The arguments are the same as for the Popen constructor. Example: + + >>> check_output(["ls", "-l", "/dev/null"]) + 'crw-rw-rw- 1 root root 1, 3 Oct 18 2007 /dev/null\n' + + The stdout argument is not allowed as it is used internally. + To capture standard error in the result, use stderr=STDOUT. + + >>> check_output(["/bin/sh", "-c", + ... "ls -l non_existent_file ; exit 0"], + ... stderr=STDOUT) + 'ls: non_existent_file: No such file or directory\n' + """ + if 'stdout' in kwargs: + raise ValueError('stdout argument not allowed, it will be overridden.') + process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs) + output, unused_err = process.communicate() + retcode = process.poll() + if retcode: + cmd = kwargs.get("args") + if cmd is None: + cmd = popenargs[0] + error = subprocess.CalledProcessError(retcode, cmd) + error.output = output + raise error + return output def debug(message): print >> sys.stderr, "[DEBUG] ", message