source: src/tests/pybin/tools.py @ fdfced6

from __future__ import print_function

import __main__
import argparse
import multiprocessing
import os
import re
import signal
import stat
import sys

from pybin import settings
from subprocess import Popen, PIPE, STDOUT

################################################################################
#               shell helpers
################################################################################

# helper functions to run terminal commands
def sh(cmd, print2stdout = True, input = None):
        # add input redirection if needed
        if input and os.path.isfile(input):
                cmd += " < %s" % input

        # if this is a dry_run, only print the commands that would be ran
        if settings.dry_run :
                print("cmd: %s" % cmd)
                return 0, None

        # otherwise create a pipe and run the desired command
        else :
                proc = Popen(cmd, stdout=None if print2stdout else PIPE, stderr=STDOUT, shell=True)
                out, err = proc.communicate()
                return proc.returncode, out

# Remove 1 or more files silently
def rm( files ):
        try:
                for file in files:
                        sh("rm -f %s > /dev/null 2>&1" % file )
        except TypeError:
                sh("rm -f %s > /dev/null 2>&1" % files )

def chdir( dest = __main__.__file__ ):
        abspath = os.path.abspath(dest)
        dname = os.path.dirname(abspath)
        os.chdir(dname)

# diff two files
def diff( lhs, rhs ):
        # diff the output of the files
        diff_cmd = ("diff --ignore-all-space "
                                "--ignore-blank-lines "
                                "--old-group-format='\t\tmissing lines :\n"
                                "%%<' \\\n"
                                "--new-group-format='\t\tnew lines :\n"
                                "%%>' \\\n"
                                "--unchanged-group-format='%%=' \\"
                                "--changed-group-format='\t\texpected :\n"
                                "%%<"
                                "\t\tgot :\n"
                                "%%>\n' \\\n"
                                "--new-line-format='\t\t%%dn\t%%L' \\\n"
                                "--old-line-format='\t\t%%dn\t%%L' \\\n"
                                "--unchanged-line-format='' \\\n"
                                "%s %s")

        # fetch return code and error from the diff command
        return sh(diff_cmd % (lhs, rhs), False)

# call make
def make(target, flags = '', redirects = '', error_file = None, silent = False):
        test_param = """test="%s" """ % (error_file) if error_file else ''
        cmd = ' '.join([
                settings.make,
                '-s' if silent else '',
                test_param,
                settings.debug.flags,
                flags,
                target,
                redirects
        ])
        return sh(cmd)

################################################################################
#               file handling
################################################################################

# helper function to replace patterns in a file
def file_replace(fname, pat, s_after):
    # first, see if the pattern is even in the file.
    with open(fname) as f:
        if not any(re.search(pat, line) for line in f):
            return # pattern does not occur in file so we are done.

    # pattern is in the file, so perform replace operation.
    with open(fname) as f:
        out_fname = fname + ".tmp"
        out = open(out_fname, "w")
        for line in f:
            out.write(re.sub(pat, s_after, line))
        out.close()
        os.rename(out_fname, fname)

# helper function to check if a files contains only a specific string
def fileContainsOnly(file, text) :
        with open(file) as f:
                ff = f.read().strip()
                result = ff == text.strip()

                return result;

# check whether or not a file is executable
def fileIsExecutable(file) :
        try :
                fileinfo = os.stat(file)
                return bool(fileinfo.st_mode & stat.S_IXUSR)
        except Exception as inst:
                print(type(inst))    # the exception instance
                print(inst.args)     # arguments stored in .args
                print(inst)
                return False

# transform path to canonical form
def canonicalPath(path):
        return os.path.join('.', os.path.normpath(path) )

# compare path even if form is different
def pathCmp(lhs, rhs):
        return canonicalPath( lhs ) == canonicalPath( rhs )

# walk all files in a path
def pathWalk( op ):
        def step(_, dirname, names):
                for name in names:
                        path = os.path.join(dirname, name)

                        op( path )

        # Start the walk
        os.path.walk('.', step, '')

################################################################################
#               system
################################################################################

# parses the Makefile to find the machine type (32-bit / 64-bit)
def getMachineType():
        sh('echo "void ?{}(int&a,int b){}int main(){return 0;}" > .dummy.c')
        ret, out = make('.dummy', silent = True)

        if ret != 0:
                print("Failed to identify architecture:")
                print(out)
                print("Stopping")
                rm( (".dummy.c",".dummy") )
                sys.exit(1)

        _, out = sh("file .dummy", print2stdout=False)
        rm( (".dummy.c",".dummy") )

        if settings.dry_run :
                return 'x64'

        return re.search(r"[^,]+,([^,]+),", out).group(1).strip()

# count number of jobs to create
def jobCount( options, tests ):
        # check if the user already passed in a number of jobs for multi-threading
        make_flags = os.environ.get('MAKEFLAGS')
        make_jobs_fds = re.search("--jobserver-(auth|fds)=\s*([0-9]+),([0-9]+)", make_flags) if make_flags else None
        if make_jobs_fds :
                tokens = os.read(int(make_jobs_fds.group(2)), 1024)
                options.jobs = len(tokens)
                os.write(int(make_jobs_fds.group(3)), tokens)
        else :
                options.jobs = multiprocessing.cpu_count()

        # make sure we have a valid number of jobs that corresponds to user input
        if options.jobs <= 0 :
                print('ERROR: Invalid number of jobs', file=sys.stderr)
                sys.exit(1)

        return min( options.jobs, len(tests) ), True if make_flags else False

# setup a proper processor pool with correct signal handling
def setupPool(jobs):
        original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
        pool = multiprocessing.Pool(jobs)
        signal.signal(signal.SIGINT, original_sigint_handler)

        return pool

# handle signals in scope
class SignalHandling():
        def __enter__(self):
                # enable signal handling
                signal.signal(signal.SIGINT, signal.SIG_DFL)

        def __exit__(self, type, value, traceback):
                # disable signal handling
                signal.signal(signal.SIGINT, signal.SIG_IGN)

################################################################################
#               misc
################################################################################

# check if arguments is yes or no
def yes_no(string):
        if string == "yes" :
                return True
        if string == "no" :
                return False
        raise argparse.ArgumentTypeError(msg)
        return False


settings.set_machine_default( getMachineType )