random_gen.py

import xml.etree.ElementTree as ElementTree
import sys
import time
import os
import random
import string
import ast
import codegen
import tchecker
import threading
from collections import deque

MAX_LENGTH = 2 ** 5 - 1


class SigFinish(Exception):
    pass


def throw_signal_function(frame, event, arg):
    raise SigFinish()


def do_nothing_trace_function(frame, event, arg):
    # Note: each function called will actually call this function
    # so, take care, your program will run slower because of that.
    return None


def interrupt_thread(thread):
    for thread_id, frame in sys._current_frames().items():
        if thread_id == thread.ident:  # Note: Python 2.6 onwards
            set_trace_for_frame_and_parents(frame, throw_signal_function)


def set_trace_for_frame_and_parents(frame, trace_func):
    # Note: this only really works if there's a tracing function set in this
    # thread (i.e.: sys.settrace or threading.settrace must have set the
    # function before)
    while frame:
        if frame.f_trace is None:
            frame.f_trace = trace_func
        frame = frame.f_back
    del frame


call_finished = False


class MyThread(threading.Thread):
    def run(self):
        # Note: this is important: we have to set the tracing function
        # when the thread is started (we could set threading.settrace
        # before starting this thread to do this externally)
        try:
            global call_finished
            sys.settrace(do_nothing_trace_function)
            exec self._Thread__kwargs['call_string'] in self._Thread__kwargs['kw']
            print "exec call finished!"
            call_finished = True
        except TypeError as t:
            call_finished = True
            self._Thread__kwargs['Exception'] = t
        except AttributeError as a:
            call_finished = True
            self._Thread__kwargs['Exception'] = a
        except SyntaxError as s1:
            call_finished = True
            self._Thread__kwargs['Exception'] = s1
        except SigFinish as sf1:
            call_finished = True
            self._Thread__kwargs['Exception'] = sf1
        except Exception as e1:
            call_finished = True
            self._Thread__kwargs['Exception'] = e1

    def __bootstrap(self):
        try:
            super(MyThread, self).__bootstrap()
        except SigFinish:
            print "SigFinish Caught - Ignore as nothing to cleanup"


class StaticAnalyzer:
    """
    This class abstracts parsing of xml file generated by Pyntch analysis.
    """
    _parsed_xml = None
    _ast_root = None

    def __init__(self, subject_file):
        """
        @subject : The name of python file to be tested. Assume to be in python path.
        """
        output_xml = "output.xml"
        parse_command = ["tchecker.py", "-o", output_xml, subject_file]
        tchecker.main(parse_command)
        self._parsed_xml = ElementTree.parse(output_xml)
        self._ast_root = ast.parse(open(subject_file).read())
        os.remove(output_xml)

    def get_functions(self):
        """
        @return : Return a list of tuples (class_node, function_node).
                  class_node is None when it is a non class function
        """
        #static functions
        all_functions = []
        for function in self._parsed_xml.findall("module/function"):
            if function.get("name").startswith("_"):
                continue
            all_functions.append((None, function))
        #end for
        #class functions
        classes = self._parsed_xml.iter("class")
        for clazz in classes:
            functions = clazz.findall("function")
            for function in functions:
                if function.get("name").startswith("_"):
                    continue
                all_functions.append((clazz, function))
                #end for
        #end for
        return all_functions

    def get_constructors(self, class_name):
        """
        @class_name : name of a class in subject
        @return : None if no such class else list of __init__ nodes.
        """
        c = None
        for clazz in self._parsed_xml.iter("class"):
            if not (clazz.get("name") == class_name):
                continue
            if c is None:
                c = []
            for function in clazz.findall("function"):
                if function.get("name") == "__init__":
                    c.append(function)
                    #endfor
                    #endfor
        return c

    def get_all_constants(self):
        """
        @return : Return constants in code as (<type>, repr(constant))
        """
        discovered_constants = dict()

        class ConstantVisitor(ast.NodeVisitor):
            def _add_constant_to_dictionary(self, ast_node):
                ast_node_type = StaticAnalyzer.get_type_from_ast_node(ast_node)
                if not (ast_node_type in discovered_constants):
                    discovered_constants[ast_node_type] = []
                discovered_constants[ast_node_type].append(codegen.to_source(ast_node))

            def visit_Dict(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_Num(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_List(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_Tuple(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_Str(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_Set(self, ast_node):
                self._add_constant_to_dictionary(ast_node)

            def visit_Call(self, ast_node):
                #print dir(node)
                func_args = ast_node.args
                #print dir(node.func)
                #print "type", type(node.func)
                if isinstance(ast_node.func, ast.Name):
                    func_name = ast_node.func.id
                    if func_name == "set" and len(func_args) == 1 and isinstance(func_args[0], ast.List):
                        (literally_list, list_elt_type) = StaticAnalyzer.get_type_from_ast_node(func_args[0])
                        ast_node_type = ("set", list_elt_type)
                        if ast_node_type in discovered_constants:
                            discovered_constants[ast_node_type].append(codegen.to_source(ast_node))
                        else:
                            discovered_constants[ast_node_type] = []
                    else:
                        self.generic_visit(ast_node)
                else:
                    self.generic_visit(ast_node)

        v = ConstantVisitor()
        v.visit(self._ast_root)
        return discovered_constants

    @staticmethod
    def get_type_from_ast_node(ast_node):
        """
        @ast_node : ast node as returned from "walk/visit"int the ast tree
        @return : type hash as would be returned by get_type_hash(xml_node) function
        """
        if isinstance(ast_node, ast.Num):
            num_literal = ast.literal_eval(codegen.to_source(ast_node))
            if isinstance(num_literal, int):
                return "int"
            else:
                return "float"
        elif isinstance(ast_node, ast.Str):
            return "str"
        elif isinstance(ast_node, ast.Dict):
            key_value = []
            if len(ast_node.keys) > 0:
                key_value.append(StaticAnalyzer.get_type_from_ast_node(ast_node.keys[0]))
            if len(ast_node.values) > 0:
                key_value.append(StaticAnalyzer.get_type_from_ast_node(ast_node.values[0]))
            if len(key_value) > 0:
                return "dict", tuple(key_value)
            else:
                return "dict"
        elif isinstance(ast_node, ast.List):
            if len(ast_node.elts) > 0:
                return "list", StaticAnalyzer.get_type_from_ast_node(ast_node.elts[0])
            else:
                return "list"
        elif isinstance(ast_node, ast.Set):
            if len(ast_node.elts) > 0:
                return "set", StaticAnalyzer.get_type_from_ast_node(ast_node.elts[0])
            else:
                return "set"
        elif isinstance(ast_node, ast.Tuple):
            components = []
            for i in ast_node.elts:
                components.append(StaticAnalyzer.get_type_from_ast_node(i))
            if len(components) > 0:
                return "tuple", tuple(components)
            else:
                return "tuple"

    @staticmethod
    def get_type_from_xml_node(node):
        """
            @node : node obtained from pyntch output xml
            @return : tuple representing the type of node
                for class ("class", <class_name>)
                for list or set ("list" or "set", <type>)
                for tuple ("tuple", <tuple_of_types>)
                for dict ("dict", <tuple_of_key_type_and_value_type>)
                for str or int or float("str" or "int" or "float", None)
                None if fail to infer type
            """
        if node is None:
            return None
        if node.tag == "compound":
            children = node.findall("./")
            if len(children) > 0:
                return StaticAnalyzer.get_type_from_xml_node(children[0])
            else:
                return None
        elif node.tag == "arg" or node.tag == "return":
            return StaticAnalyzer.get_type_from_xml_node(node.find("compound"))
        elif node.tag == "class":
            return "class", node.get("name")
        elif node.tag == "int" or node.tag == "float" or node.tag == "str":
            return node.tag
        elif node.tag == "list" or node.tag == "set":
            compounds = node.findall("compound")
            if len(compounds) > 0:
                return node.tag, StaticAnalyzer.get_type_from_xml_node(compounds[0])
            else:
                return node.tag
        elif node.tag == "dict":
            compounds = node.findall("compound")
            if len(compounds) > 0:
                if len(compounds) == 1:
                    compounds.append(compounds[0])
                key_value = [StaticAnalyzer.get_type_from_xml_node(compounds[0]),
                             StaticAnalyzer.get_type_from_xml_node(compounds[1])]
                return "dict", tuple(key_value)
            else:
                return "dict"
        elif node.tag == "tuple":
            type_list = []
            for component in node.findall("./"):
                type_list.append(component.tag)
            if len(type_list) > 0:
                return "tuple", tuple(type_list)
            else:
                return "tuple"
        elif node.tag == "instance":
            moduled_name = node.get("name").split(".")
            if len(moduled_name) >= 2:
                return "class", node.get("name").split(".")[1]
            elif len(moduled_name) == 1:
                return "class", node.get("name").split(".")[0]
            else:
                return None

        else:
            return None


class SequenceGenerator:
    """
    This class abstracts the generation of function call sequence.
    """
    _max_size = None
    _function_nodes = None
    _subject = None
    _analyzer = None

    def __init__(self, ms, subject_file):
        """
        @max_size : The maximum size of the generated sequence.
        @subject : The name of python file to be tested. Assume to be in pythonpath.
        """
        self._max_size = ms
        self._analyzer = StaticAnalyzer(subject_file)
        self._function_nodes = self._analyzer.get_functions()
        self._subject = os.path.splitext(os.path.basename(subject_file))[0]

    def _generate_arg(self, arg_node, previous_results, static_constants, max_values):
        """
        @arg_node : A xml node of type arg that belongs to a function.
        @previous_results : A dictionary containing results of previous calls. Key : type, Value : result_variable
        @static_constants : A dictionary containing constants found in code. Key : type, Value : list of const. literals
        @max_values: Dictionary containing the max_size of value to be generated. Key: "str", "int", "list", Value : Int
        @return : Serialized value to be passed a function call. <name>=<value>
        """
        #print "generating arg for " + arg_node.get('name')
        arg_type = StaticAnalyzer.get_type_from_xml_node(arg_node)
        name = arg_node.get('name')
        serialized_arg = name + "="
        if arg_type in previous_results:
            return serialized_arg + random.choice(previous_results[arg_type])
        compounds = arg_node.findall("./compound")
        for compound in compounds:
            type_nodes = compound.findall("./")
            type_node = None
            if len(type_nodes) > 0:
                type_node = type_nodes[0]
            serialized_arg = serialized_arg + self._generate_value(type_node, previous_results, static_constants,
                                                                   max_values)
        if serialized_arg[-1] == "=":
            serialized_arg += "None"
        return serialized_arg

    def _generate_value(self, node, previous_results, static_constants, value_hints):
        """
        @node : A xml node of type int, str,...
        @previous_results : A dictionary containing results of previous calls. Key : type, Value : result_variable
        @static_constants : A dictionary containing constants found in code. Key : type, Value : list of const. literals
        @max_values: Dictionary containing the max_size of value to be generated. Key: "str", "int", "list", Value : Int
        @return : A random value of type indicated by node.
        """
        #print "generating value for " + node.tag
        node_type = StaticAnalyzer.get_type_from_xml_node(node)
        if node_type is None:
            if random.choice([1, 2, 3, 4]) == 4:
                return repr(None)
            else:
                if len(previous_results) > 0:
                    choosers = random.choice(previous_results.values())
                    return random.choice(choosers)
                else:
                    return repr(None)
            #print "Existing results", previous_results
            #choosers = random.choice(previous_results.values() + [deque([repr(None)])])
            #return random.choice(choosers)
        if (node_type in static_constants) and (random.uniform(0, 1) < 0.25) and value_hints['use_static_constants']:
            return random.choice(static_constants.get(node_type))
        if node_type in previous_results:
            return random.choice(previous_results[node_type])
        if node.tag == "class" or node.tag == "instance":
            if node.tag == "instance":
                class_name = node.get("name").split(".")[1]
            else:
                class_name = node.get("name")
            constructors = self._analyzer.get_constructors(class_name)
            if len(constructors) == 0:
                return self._subject + "." + class_name + "()"
            constructor = constructors[random.randint(0, len(constructors) - 1)]
            value = self._subject + "." + class_name + "( "
            args = constructor.findall("arg")
            for arg in args:
                if arg.get("name") == "self":
                    continue
                value += self._generate_arg(arg, previous_results, static_constants, value_hints) + ","
            #end for
            value = value[:-1] + ")"
            return value
        if node.tag == "compound":
            return self._generate_value(node[random.randint(0, len(node.findall("./")) - 1)], previous_results,
                                        static_constants, value_hints)
        if node.tag == "int":
            if value_hints['int'] == "+ve":
                return repr(random.randint(0, MAX_LENGTH))
            else:
                return repr(random.randint(-1 * MAX_LENGTH, MAX_LENGTH))
        if node.tag == "float":
            if value_hints['int'] == "+ve":
                return repr(random.uniform(0, MAX_LENGTH))
            else:
                return repr(random.uniform(-1 * MAX_LENGTH, MAX_LENGTH))
        if node.tag == "str":
            v = ''
            str_length = random.randint(0, MAX_LENGTH)
            for i in range(str_length):
                if value_hints['str'] == 'asc':
                    v += random.choice(string.printable)
                else:
                    v += chr(random.randint(0, 255))
            #endfor
            return repr(v)
        if node.tag == 'tuple':
            #max_values['str'] = 2 ** 5 - 1
            v = '( '
            for component in node:
                v += self._generate_value(component, previous_results, static_constants, value_hints) + ","
            v = v[:-1] + ")"
            return v
        if node.tag == "list" or node.tag == "set":
            #max_values['str'] = 2 ** 5 - 1
            v = ""
            compounds = node.findall("./compound")
            for i in range(0, random.randint(0, value_hints['list'])):
                for compound in compounds:
                    if v != "":
                        v += ", "
                    v = v + self._generate_value(compound, previous_results, static_constants, value_hints)
            v = "[" + v + "]"
            if node.tag == "set":
                v = "set(" + v + ")"
            #max_values['list'] = 2 ** 5 - 1
            return v
        if node.tag == 'dict':
            #max_values['str'] = 2 ** 5 - 1
            compounds = node.findall("./compound")
            if len(compounds) < 1:
                return "{}"
            v = "{ "
            if len(compounds) == 1:
                compounds.append(compounds[0])
            for i in range(0, random.randint(0, value_hints['dict'])):
                v += self._generate_value(compounds[0], previous_results, static_constants,
                                          value_hints) + ":" + self._generate_value(compounds[1], previous_results,
                                                                                    static_constants, value_hints) + ","
            v = v[:-1] + "}"
            #max_values['dict'] = 2 ** 5 - 1
            return v
        return "None"

    def _add_to_results(self, results, type_key, repr_value):
        if type_key is None:
            return
        if not (type_key in results):
            results[type_key] = deque()
        results[type_key].append(repr_value)
        if len(results[type_key]) > 5:
            results[type_key].popleft()

    def get_random_sequence(self, use_previous, time_remaining, **kw):
        """
        @use_previous : If True, we use previous calls in sequence
        @timeout : maximum time in seconds of a call
        @kw : import namespace to use for checking if the calls generated are valid.
        @return : List of serialized function calls to be inserted in TestCase file.
        """
        generation_start_time = time.time()
        calls = []
        total_functions = len(self._function_nodes)
        base = random.randint(0, total_functions - 1)
        var_id = 0
        results = dict()
        static_constants = self._analyzer.get_all_constants()
        value_hints = {'list': MAX_LENGTH, 'dict': MAX_LENGTH, 'int': random.choice(['+ve', 'mix']),
                       'str': random.choice(['asc', 'txt']), 'use_static_constants': random.choice([True, False])}

        for i in range(self._max_size):
            if (time.time() - generation_start_time) >= time_remaining:
                break
            #max_values = {'str': 2 ** 10 - 1, 'list': 2 ** 10 - 1, 'dict': 2 ** 10 - 1,
            (clazz, f) = self._function_nodes[(base + i) % total_functions]
            function_return_var = "return" + str(var_id)
            var_id += 1
            if clazz is None:
                call = function_return_var + " = " + self._subject + "." + f.get("name") + "("
            else:
                constructor_call = self._generate_value(clazz, results, static_constants, value_hints)
                if constructor_call is None:
                    continue
                class_return_var = "return" + str(var_id)
                var_id += 1
                call = class_return_var + " = " + constructor_call + "\n"
                call += function_return_var + " = " + class_return_var + "." + f.get("name") + "("
            arg_list = ""
            arg_nodes = f.findall("./arg")
            arg_dict = dict()
            for arg_node in arg_nodes:
                if not (clazz is None) and arg_node.get('name') == "self":
                    continue
                if arg_list != "":
                    arg_list += ", "
                serialized_arg = self._generate_arg(arg_node, results, static_constants, value_hints)
                if use_previous and len(serialized_arg.split("=")) > 0:
                    serialized_arg_type = StaticAnalyzer.get_type_from_xml_node(arg_node)
                    self._add_to_results(arg_dict, serialized_arg_type, serialized_arg[serialized_arg.find("=") + 1:])
                    print "Adding serialized arg to results", serialized_arg_type, \
                        serialized_arg[serialized_arg.find("=") + 1:]
                arg_list = arg_list + serialized_arg
            call += arg_list + ")"
            print "Running " + call,
            global call_finished
            call_finished = False
            args = {'call_string': call, 'kw': kw, 'Exception': None}
            #thread = threading.Thread(target=calling_thread, args=[call, kw])
            thread = MyThread(kwargs=args)
            thread.start()
            time.sleep(0.1)  # Warm up time
            TIMEOUT_PER_CALL = 5
            elapsed = 0
            print "Checking for timeout(" + str(TIMEOUT_PER_CALL) + ")",
            while (not call_finished) and elapsed < TIMEOUT_PER_CALL:
                if elapsed % 4 == 0:
                    print ".",
                time.sleep(0.25)
                elapsed += 0.25
            if not call_finished:
                print "\nInterrupting Thread-"
                interrupt_thread(thread)
            print 'Joining'
            thread.join()  # Joining here: if we didn't interrupt the thread before, we'd be here forever.
            print 'Finished'
            exec_exception = args['Exception']
            if exec_exception is None:
                if use_previous:
                    if not (clazz is None) and not (class_return_var is None):
                        class_type = StaticAnalyzer.get_type_from_xml_node(clazz)
                        self._add_to_results(results, class_type, class_return_var)
                    type_hash = StaticAnalyzer.get_type_from_xml_node(f.find("return"))
                    self._add_to_results(results, type_hash, function_return_var)
                    for arg_type_key, serialized_arg_values_deque in arg_dict.iteritems():
                        for deque_item in serialized_arg_values_deque:
                            self._add_to_results(results, arg_type_key, deque_item)
                print "PASS - Add to sequence"
                calls += call.split("\n")
            elif isinstance(exec_exception, TypeError):
                print "FATAL - TypeError, Discarding", exec_exception
            elif isinstance(exec_exception, AttributeError):
                print "FATAL - AttributeError, Discarding", exec_exception
            elif isinstance(exec_exception, SyntaxError):
                print "FATAL - Invalid Syntax, Probably Bug, Discarding", exec_exception
            elif isinstance(exec_exception, SigFinish):
                print "FATAL - Too long call, Discarding", SigFinish
            else:
                if i > 0.7 * self._max_size:
                    print "PASS LAST - Unknown exception, Terminating sequence"
                    calls += call.split("\n")
                    break
                else:
                    print "FATAL - Unknown exception, Discarding and Continuing."

        return calls


class TestWriter:
    """
    Helper class for generating the final test file. The test cases are output
    in accordance with the TestCase framework.
    """
    _file = None
    _test_num = None
    _indent = None
    _subject = None

    def _get_base_code(self):
        return '''\
import {0}
import unittest

class RandomTestCases(unittest.TestCase) :
'''.format(self._subject)

    def __init__(self, od, subject_module):
        """
        @output_dir : The output directory where test cases should be written.
        @subject : The subject filename to be imported in test file.
        """
        file_name = "Test_" + subject_module + "_Test.py"
        self._subject = subject_module
        self._file = open(os.path.join(od, file_name), 'w')
        self._file.write(self._get_base_code())
        self._test_num = 0
        self._indent = 2

    def _get_test_function(self, call_sequence):
        """
        @call_sequence : list of serialized function calls.
        @return : serialized test function containing those calls.
        """
        joiner = '\n' + ' ' * self._indent
        serialized_function = ''
        base_string = '''\
def test{0}(self) :
{1}
'''.format(str(self._test_num), joiner + joiner.join(call_sequence))
        lines = base_string.split('\n')
        for line in lines:
            serialized_function += ' ' * self._indent + line + '\n'
        return serialized_function

    def serialize_to_test(self, call_sequence):
        """
        @call_sequence : list of serialized function calls.
        @return : None, Makes a test function from the call_sequence. Writes it to
                  test file.
        """
        if len(call_sequence) > 0:
            self._file.write(self._get_test_function(call_sequence))
            self._test_num += 1

    def close(self):
        """
        @return : None, Creates the final and valid PyUnit test file.
        """
        if self._test_num == 0:
            self._file.write(' ' * self._indent + "pass\n")

        end_string = '''\
if __name__ == '__main__':
unittest.main()
suite = unittest.TestLoader().loadTestsFromTestCase(RandomTestCases)
unittest.TextTestRunner(verbosity=2).run(suite)
'''
        lines = end_string.split('\n')
        joiner = '\n' + ' ' * self._indent
        fixed_string = joiner.join(lines)
        self._file.write(fixed_string)
        self._file.close()


if __name__ == "__main__":
    start = current = time.time()
    random.seed()
    if len(sys.argv) < 5:
        print "Usage : random_gen.py <subject> <max_size> <timeout> <output_dir>"
        sys.exit(0)
    subject = sys.argv[1]
    max_size = int(sys.argv[2])
    if max_size == 0:
        max_size = 1
    timeout = int(sys.argv[3]) - 3  # Be safe here.
    output_dir = sys.argv[4]
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    sequence_generator = SequenceGenerator(max_size, subject)
    subject_name = os.path.splitext(os.path.basename(subject))[0]
    test_writer = TestWriter(output_dir, subject_name)
    while int(current - start) < timeout:
        try:
            current = time.time()
            s = sequence_generator.get_random_sequence(random.choice([True, False]),
                                                       timeout - int(current - start),
                                                       **{subject_name: __import__(subject_name)})
            test_writer.serialize_to_test(s)
        except SigFinish as sf:
            print "Caught sigfinish"
        except Exception as e:
            print "error generating sequence, discarding, trying next one...", e
    test_writer.close()
    sys.exit(0)