def test_count_all_messages(self): # SETUP # EXEC builder = ReportBuilder() frame = builder.start_frame(1, 2) frame.add_message('a', 1) frame.add_message('b', 2) # VERIFY self.assertEqual(0, builder.message_count) self.assertEqual(2, builder.count_all_messages())
def trace_code(self, source, load_as=None, is_module=False, dump=False, driver=None, filename=None, bad_driver=None): """ Trace a module of source code, possibly by running a driver script. :param str source: the source code to trace :param str load_as: the module name to load the source code as :param bool is_module: True if the driver is a module name instead of a file name :param bool dump: True if the source code should be included in the output :param list driver: the driver script's file name or module name and args :param str filename: the file name of the source code :param str bad_driver: a message to display if the driver doesn't call the module """ builder = ReportBuilder(self.message_limit) builder.max_width = self.max_width self.return_code = 0 try: tree = parse(source) new_tree = Tracer().visit(tree) fix_missing_locations(new_tree) LineNumberCleaner().visit(new_tree) # from ast import dump # print(dump(new_tree, include_attributes=True)) code = compile(new_tree, PSEUDO_FILENAME, 'exec') self.environment[CONTEXT_NAME] = builder is_own_driver = is_module and driver and driver[0] == load_as seed(0) self.run_instrumented_module(code, load_as, filename, is_own_driver) if driver and not is_own_driver: start_count = builder.message_count with self.swallow_output(): try: if not is_module: self.run_python_file(driver[0], driver) end_count = builder.count_all_messages() else: module_name = driver[0] self.run_python_module(module_name, driver) end_count = builder.count_all_messages() except SystemExit as ex: end_count = builder.count_all_messages() if ex.code: self.return_code = ex.code messages = traceback.format_exception_only( type(ex), ex) message = messages[-1].strip() self.report_driver_result(builder, [message]) if end_count == start_count: driver_name = os.path.basename(driver[0]) message = (bad_driver or "{} doesn't call the {} module." " Try a different driver.".format( driver_name, load_as)) self.report_driver_result(builder, [message]) for value in self.environment.values(): if isinstance(value, types.GeneratorType): value.close() except SyntaxError: self.return_code = 1 ex = sys.exc_info()[1] messages = traceback.format_exception_only(type(ex), ex) builder.add_message(messages[-1].strip() + ' ', ex.lineno) except BaseException as ex: self.return_code = getattr(ex, 'code', 1) etype, value, tb = sys.exc_info() is_reported = False messages = traceback.format_exception_only(etype, value) entries = traceback.extract_tb(tb) for filename, _, _, _ in entries: if filename == PSEUDO_FILENAME: is_reported = True while not is_reported and tb is not None: frame = tb.tb_frame code = frame.f_code filename = code.co_filename if __file__ not in (filename, filename + 'c'): break tb = tb.tb_next if not is_reported: if tb: messages = traceback.format_exception(etype, value, tb) else: messages = traceback.format_exception_only(etype, value) self.report_driver_result(builder, messages) report = builder.report(source.count('\n')) if dump: source_lines = source.splitlines() report_lines = report.splitlines() dump_lines = [] source_width = max(map(len, source_lines)) indent = 4 for source_line, report_line in izip_longest(source_lines, report_lines, fillvalue=''): line = (indent * ' ' + source_line + (source_width - len(source_line)) * ' ' + ' | ' + report_line) dump_lines.append(line) report = '\n'.join(dump_lines) return report