def add_test_case(test_suite, test_case): if isinstance(test_suite, Module): # If the main_snippet exists we have to put the new test case # before it. If it doesn't we put the test case at the end. main_snippet = code_of(test_suite, 'main_snippet') if main_snippet: insert_before(main_snippet, test_case.code) else: code_of(test_suite).append_child(test_case.code) elif isinstance(test_suite, TestClass): # Append to the right node, so that indentation level of the # new method is good. if code_of(test_suite).children and is_node_of_type( code_of(test_suite).children[-1], 'suite'): remove_trailing_whitespace(test_case.code) suite = code_of(test_suite).children[-1] # Prefix the definition with the right amount of whitespace. node = find_last_leaf(suite.children[-2]) ident = get_starting_whitespace(suite) # There's no need to have extra newlines. if node.prefix.endswith("\n"): node.prefix += ident.lstrip("\n") else: node.prefix += ident # Insert before the class contents dedent. suite.insert_child(-1, test_case.code) else: code_of(test_suite).append_child(test_case.code) else: raise TypeError("Tried to add a test case to %r." % test_suite) add_test_case_without_append(test_suite, test_case) test_suite.mark_as_changed()
def add_test_case(test_suite, test_case): if isinstance(test_suite, Module): # If the main_snippet exists we have to put the new test case # before it. If it doesn't we put the test case at the end. main_snippet = code_of(test_suite, 'main_snippet') if main_snippet: insert_before(main_snippet, test_case.code) else: code_of(test_suite).append_child(test_case.code) elif isinstance(test_suite, TestClass): # Append to the right node, so that indentation level of the # new method is good. if code_of(test_suite).children and is_node_of_type(code_of(test_suite).children[-1], 'suite'): remove_trailing_whitespace(test_case.code) suite = code_of(test_suite).children[-1] # Prefix the definition with the right amount of whitespace. node = find_last_leaf(suite.children[-2]) ident = get_starting_whitespace(suite) # There's no need to have extra newlines. if node.prefix.endswith("\n"): node.prefix += ident.lstrip("\n") else: node.prefix += ident # Insert before the class contents dedent. suite.insert_child(-1, test_case.code) else: code_of(test_suite).append_child(test_case.code) else: raise TypeError("Tried to add a test case to %r." % test_suite) add_test_case_without_append(test_suite, test_case) test_suite.mark_as_changed()
def parse_fragment(code): """Works like parse() but returns an object stripped of the file_input wrapper. This eases merging this piece of code into other ones. """ parsed_code = parse(code) if is_node_of_type(parsed_code, 'file_input') and \ len(parsed_code.children) == 2 and \ is_leaf_of_type(parsed_code.children[-1], token.ENDMARKER): return parsed_code.children[0] return parsed_code