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
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
