def insert_output_start_stop_indicators(src):
"""
Insert identifier strings so that output can be segregated from input.
Parameters
----------
src : str
String containing input and output lines.
Returns
-------
str
String with output demarked.
"""
rb = RedBaron(src)
# find lines with trailing comments so we can preserve them properly
lines_with_comments = {}
comments = rb.findAll('comment')
for c in comments:
if c.previous and c.previous.type != 'endl':
lines_with_comments[c.previous] = c
input_block_number = 0
# find all nodes that might produce output
nodes = rb.findAll(lambda identifier: identifier in ['print', 'atomtrailers'])
for r in nodes:
# assume that whatever is in the try block will fail and produce no output
# this way we can properly handle display of error messages in the except
if hasattr(r.parent, 'type') and r.parent.type == 'try':
continue
# Output within if/else statements is not a good idea for docs, because
# we don't know which branch execution will follow and thus where to put
# the output block. Regardless of which branch is taken, though, the
# output blocks must start with the same block number.
if hasattr(r.parent, 'type') and r.parent.type == 'if':
if_block_number = input_block_number
if hasattr(r.parent, 'type') and r.parent.type in ['elif', 'else']:
input_block_number = if_block_number
if is_output_node(r):
# if there was a trailing comment on this line, output goes after it
if r in lines_with_comments:
r = lines_with_comments[r] # r is now the comment
# find the correct node to 'insert_after'
while hasattr(r, 'parent') and not hasattr(r.parent, 'insert'):
r = r.parent
r.insert_after('print(">>>>>%d")\n' % input_block_number)
input_block_number += 1
# curse you, redbaron! stop inserting endl before trailing comments!
for l, c in lines_with_comments.items():
if c.previous and c.previous.type == 'endl':
c.previous.value = ''
return rb.dumps()
def replace_asserts_with_prints(source_code):
"""
Replace asserts with print statements.
Using RedBaron, replace some assert calls with print statements that print the actual
value given in the asserts.
Depending on the calls, the actual value can be the first or second
argument.
"""
rb = RedBaron(source_code) # convert to RedBaron internal structure
for assert_type in [
'assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
'assertEqualArrays', 'assertTrue', 'assertFalse'
]:
assert_nodes = rb.findAll("NameNode", value=assert_type)
for assert_node in assert_nodes:
assert_node = assert_node.parent
remove_redbaron_node(assert_node, 0) # remove 'self' from the call
assert_node.value[0].replace('print')
if assert_type not in ['assertTrue', 'assertFalse']:
remove_redbaron_node(assert_node.value[1],
1) # remove the expected value argument
assert_nodes = rb.findAll("NameNode", value='assert_rel_error')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 4 arguments
if len(assert_node.value[1]) == 4:
remove_redbaron_node(assert_node.value[1],
-1) # remove the relative error tolerance
remove_redbaron_node(assert_node.value[1],
-1) # remove the expected value
remove_redbaron_node(assert_node.value[1],
0) # remove the first argument which is
# the TestCase
assert_node.value[0].replace("print")
assert_nodes = rb.findAll("NameNode", value='assert_almost_equal')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 3 arguments
if len(assert_node.value[1]) == 3:
remove_redbaron_node(assert_node.value[1],
-1) # remove the relative error tolerance
remove_redbaron_node(assert_node.value[1],
-1) # remove the expected value
assert_node.value[0].replace("print")
source_code_with_prints = rb.dumps(
) # get back the string representation of the code
return source_code_with_prints
def _strip_asserts(source):
"""
Remove assert method calls from source code.
Using RedBaron, replace some assert calls with print statements that print the actual
value given in the asserts. Depending on the calls, the actual value can be the first or second
argument.
Parameters
----------
source : str
String containing source lines.
Returns
-------
str
Source with asserts removed.
"""
rb = RedBaron(source) # convert to RedBaron internal structure
# findAll is slow, so only check the ones that are present.
asserts = ['assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
'assert_equal_arrays', 'assertTrue', 'assertFalse', 'assert_near_equal',
'assert_rel_error', 'assert_almost_equal', 'assert_allclose']
for assert_type in asserts:
assert_nodes = rb.findAll("NameNode", value=assert_type)
for i in reversed(range(len(assert_nodes))):
parent = assert_nodes[i].parent
for j in reversed(range(len(parent.value))):
assert_nodes[i].parent.remove(parent.value[j])
return rb.dumps()
def replace_asserts_with_prints(source_code):
"""
Replace asserts with print statements.
Using RedBaron, replace some assert calls with print statements that print the actual
value given in the asserts.
Depending on the calls, the actual value can be the first or second
argument.
"""
rb = RedBaron(source_code) # convert to RedBaron internal structure
for assert_type in ['assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
'assert_equal_arrays', 'assertTrue', 'assertFalse']:
assert_nodes = rb.findAll("NameNode", value=assert_type)
for assert_node in assert_nodes:
assert_node = assert_node.parent
remove_redbaron_node(assert_node, 0) # remove 'self' from the call
assert_node.value[0].replace('print')
if assert_type not in ['assertTrue', 'assertFalse']:
remove_redbaron_node(assert_node.value[1], 1) # remove the expected value argument
assert_nodes = rb.findAll("NameNode", value='assert_rel_error')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 4 arguments
if len(assert_node.value[1]) == 4:
remove_redbaron_node(assert_node.value[1], -1) # remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1) # remove the expected value
remove_redbaron_node(assert_node.value[1], 0) # remove the first argument which is
# the TestCase
assert_node.value[0].replace("print")
assert_nodes = rb.findAll("NameNode", value='assert_almost_equal')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 3 arguments
if len(assert_node.value[1]) == 3:
remove_redbaron_node(assert_node.value[1], -1) # remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1) # remove the expected value
assert_node.value[0].replace("print")
source_code_with_prints = rb.dumps() # get back the string representation of the code
return source_code_with_prints
def remove_raise_skip_tests(src):
"""
Remove from the code any raise unittest.SkipTest lines since we don't want those in
what the user sees.
"""
rb = RedBaron(src)
raise_nodes = rb.findAll("RaiseNode")
for rn in raise_nodes:
# only the raise for SkipTest
if rn.value[:2].dumps() == 'unittestSkipTest':
rn.parent.value.remove(rn)
return rb.dumps()
def get_method_body(method_code):
'''Using the RedBaron module, get the body of a method.
Do not want the definition signature line
'''
method_code = '\n' + method_code # For some reason RedBaron has problems with this if
# if it does not start with an empty line
rb = RedBaron(method_code)
def_node = rb.findAll("DefNode")[0] # Look for the 'def' node. Should only be one!
def_node.value.decrease_indentation(8) # so that the code is all the way to the left
return def_node.value.dumps()
def remove_raise_skip_tests(src):
"""
Remove from the code any raise unittest.SkipTest lines since we don't want those in
what the user sees.
"""
rb = RedBaron(src)
raise_nodes = rb.findAll("RaiseNode")
for rn in raise_nodes:
# only the raise for SkipTest
if rn.value[:2].dumps() == 'unittestSkipTest':
rn.parent.value.remove(rn)
return rb.dumps()
def get_method_body(method_code):
'''Using the RedBaron module, get the body of a method.
Do not want the definition signature line
'''
method_code = '\n' + method_code # For some reason RedBaron has problems with this if
# if it does not start with an empty line
rb = RedBaron(method_code)
def_node = rb.findAll("DefNode")[0] # Look for the 'def' node. Should only be one!
def_node.value.decrease_indentation(8) # so that the code is all the way to the left
return def_node.value.dumps()
def get_skip_predicate_and_message(source, method_name):
'''
Look to see if the method has a unittest.skipUnless or unittest.skip
decorator.
If it has a unittest.skipUnless decorator, return the predicate and the message
If it has a unittest.skip decorator, return just the message ( set predicate to None )
'''
rb = RedBaron(source)
def_nodes = rb.findAll("DefNode", name=method_name)
if def_nodes:
if def_nodes[0].decorators:
if def_nodes[0].decorators[0].value.dumps(
) == 'unittest.skipUnless':
return (
def_nodes[0].decorators[0].call.value[0].dumps(),
def_nodes[0].decorators[0].call.value[1].value.to_python())
elif def_nodes[0].decorators[0].value.dumps() == 'unittest.skip':
return (
None,
def_nodes[0].decorators[0].call.value[0].value.to_python())
return None
def replace_asserts_with_prints(src):
"""
Replace asserts with print statements.
Using RedBaron, replace some assert calls with print statements that print the actual
value given in the asserts. Depending on the calls, the actual value can be the first or second
argument.
Parameters
----------
src : str
String containing source lines.
Returns
-------
str
String containing source with asserts replaced by prints.
"""
rb = RedBaron(src) # convert to RedBaron internal structure
# findAll is slow, so only check the ones that are present.
base_assert = [
'assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
'assert_equal_arrays', 'assertTrue', 'assertFalse'
]
used_assert = [item for item in base_assert if item in src]
for assert_type in used_assert:
assert_nodes = rb.findAll("NameNode", value=assert_type)
for assert_node in assert_nodes:
assert_node = assert_node.parent
remove_redbaron_node(assert_node, 0) # remove 'self' from the call
assert_node.value[0].replace('print')
if assert_type not in ['assertTrue', 'assertFalse']:
# remove the expected value argument
remove_redbaron_node(assert_node.value[1], 1)
if 'assert_rel_error' in src:
assert_nodes = rb.findAll("NameNode", value='assert_rel_error')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 4 arguments
if len(assert_node.value[1]) == 4:
# remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1)
remove_redbaron_node(assert_node.value[1],
-1) # remove the expected value
# remove the first argument which is the TestCase
remove_redbaron_node(assert_node.value[1], 0)
#
assert_node.value[0].replace("print")
if 'assert_almost_equal' in src:
assert_nodes = rb.findAll("NameNode", value='assert_almost_equal')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 3 arguments
if len(assert_node.value[1]) == 3:
# remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1)
remove_redbaron_node(assert_node.value[1],
-1) # remove the expected value
assert_node.value[0].replace("print")
return rb.dumps()
def replace_asserts_with_prints(src):
"""
Replace asserts with print statements.
Using RedBaron, replace some assert calls with print statements that print the actual
value given in the asserts. Depending on the calls, the actual value can be the first or second
argument.
Parameters
----------
src : str
String containing source lines.
Returns
-------
str
String containing source with asserts replaced by prints.
"""
rb = RedBaron(src) # convert to RedBaron internal structure
# findAll is slow, so only check the ones that are present.
base_assert = ['assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
'assert_equal_arrays', 'assertTrue', 'assertFalse']
used_assert = [item for item in base_assert if item in src]
for assert_type in used_assert:
assert_nodes = rb.findAll("NameNode", value=assert_type)
for assert_node in assert_nodes:
assert_node = assert_node.parent
remove_redbaron_node(assert_node, 0) # remove 'self' from the call
assert_node.value[0].replace('print')
if assert_type not in ['assertTrue', 'assertFalse']:
# remove the expected value argument
remove_redbaron_node(assert_node.value[1], 1)
if 'assert_rel_error' in src:
assert_nodes = rb.findAll("NameNode", value='assert_rel_error')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 4 arguments
if len(assert_node.value[1]) == 4:
# remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1)
remove_redbaron_node(assert_node.value[1], -1) # remove the expected value
# remove the first argument which is the TestCase
remove_redbaron_node(assert_node.value[1], 0)
#
assert_node.value[0].replace("print")
if 'assert_almost_equal' in src:
assert_nodes = rb.findAll("NameNode", value='assert_almost_equal')
for assert_node in assert_nodes:
assert_node = assert_node.parent
# If relative error tolerance is specified, there are 3 arguments
if len(assert_node.value[1]) == 3:
# remove the relative error tolerance
remove_redbaron_node(assert_node.value[1], -1)
remove_redbaron_node(assert_node.value[1], -1) # remove the expected value
assert_node.value[0].replace("print")
return rb.dumps()
def get_unit_test_source_and_run_outputs(method_path):
'''
Get the source code for a unit test method, run the test,
and capture the output of the run
'''
module_path = '.'.join(method_path.split('.')[:-2])
class_name = method_path.split('.')[-2]
method_name = method_path.split('.')[-1]
test_module = importlib.import_module(module_path)
cls = getattr(test_module, class_name)
try:
import mpi4py
except ImportError:
use_mpi = False
else:
N_PROCS = getattr(cls, 'N_PROCS', 1)
use_mpi = N_PROCS > 1
meth = getattr(cls, method_name)
class_source_code = inspect.getsource(cls)
rb = RedBaron(class_source_code)
def_nodes = rb.findAll("DefNode", name=method_name)
def_nodes[0].value.decrease_indentation(8)
method_source = def_nodes[0].value.dumps()
# Remove docstring from source code
source_minus_docstrings = remove_docstrings(method_source)
# We are using the RedBaron module in the next two function calls
# to get the code in the way we want it.
# Only want the method body. Do not want the 'def' line
# method_body_source = get_method_body(source_minus_docstrings)
method_body_source = source_minus_docstrings
# Replace some of the asserts with prints of the actual values
source_minus_docstrings_with_prints = replace_asserts_with_prints(
method_body_source)
# remove raise SkipTest lines
# We decided to leave them in for now
# source_minus_docstrings_with_prints = remove_raise_skip_tests(source_minus_docstrings_with_prints)
# Remove the initial empty lines
source_minus_docstrings_with_prints_cleaned = remove_initial_empty_lines_from_source(
source_minus_docstrings_with_prints)
# Get all the pieces of code needed to run the unit test method
module_source_code = inspect.getsource(test_module)
lines_before_test_cases = get_lines_before_test_cases(module_source_code)
setup_source_code = get_method_body(
inspect.getsource(getattr(cls, 'setUp')))
teardown_source_code = get_method_body(
inspect.getsource(getattr(cls, 'tearDown')))
# If the test method has a skipUnless or skip decorator, we need to convert it to a
# raise call
skip_predicate_and_message = get_skip_predicate_and_message(
class_source_code, method_name)
if skip_predicate_and_message:
# predicate, message = skip_unless_predicate_and_message
predicate, message = skip_predicate_and_message
if predicate:
raise_skip_test_source_code = 'import unittest\nif not {}: raise unittest.SkipTest("{}")'.format(
predicate, message)
else:
raise_skip_test_source_code = 'import unittest\nraise unittest.SkipTest("{}")'.format(
message)
else:
raise_skip_test_source_code = ""
code_to_run = '\n'.join([
lines_before_test_cases, setup_source_code,
raise_skip_test_source_code,
source_minus_docstrings_with_prints_cleaned, teardown_source_code
])
# Write it to a file so we can run it. Tried using exec but ran into problems with that
fd, code_to_run_path = tempfile.mkstemp()
skipped = False
failed = False
try:
with os.fdopen(fd, 'w') as tmp:
tmp.write(code_to_run)
tmp.close()
if use_mpi:
env = os.environ.copy()
env['USE_PROC_FILES'] = '1'
p = subprocess.Popen(
['mpirun', '-n',
str(N_PROCS), 'python', code_to_run_path],
env=env)
p.wait()
multi_out_blocks = []
for i in range(N_PROCS):
with open('%d.out' % i) as f:
multi_out_blocks.append(extract_output_blocks(f.read()))
os.remove('%d.out' % i)
output_blocks = []
for i in range(len(multi_out_blocks[0])):
output_blocks.append('\n'.join([
"(rank %d) %s" % (j, m[i])
for j, m in enumerate(multi_out_blocks) if m[i]
]))
else:
run_outputs = subprocess.check_output(['python', code_to_run_path],
stderr=subprocess.STDOUT)
except subprocess.CalledProcessError as e:
# Get a traceback like this:
# Traceback (most recent call last):
# File "/Applications/PyCharm CE.app/Contents/helpers/pydev/pydevd.py", line 1556, in <module>
# globals = debugger.run(setup['file'], None, None, is_module)
# File "/Applications/PyCharm CE.app/Contents/helpers/pydev/pydevd.py", line 940, in run
# pydev_imports.execfile(file, globals, locals) # execute the script
# File "/var/folders/l3/9j86k5gn6cx0_p25kdplxgpw1l9vkk/T/tmp215aM1", line 23, in <module>
# raise unittest.SkipTest("check_totals not implemented yet")
# unittest.case.SkipTest: check_totals not implemented yet
if 'raise unittest.SkipTest' in e.output.decode('utf-8'):
reason_for_skip = e.output.splitlines(
)[-1][len('unittest.case.SkipTest: '):]
run_outputs = reason_for_skip
skipped = True
else:
run_outputs = "Running of embedded test {} in docs failed due to: \n\n{}".format(
method_path, e.output.decode('utf-8'))
failed = True
# print("Running of embedded test " + method_path + " in docs failed due to: " + e.output.decode('utf-8'))
# raise
except Exception as err:
run_outputs = "Running of embedded test {} in docs failed due to: \n\n{}".format(
method_path, str(err))
failed = True
finally:
os.remove(code_to_run_path)
if PY3:
run_outputs = "".join(map(
chr, run_outputs)) # in Python 3, run_outputs is of type bytes!
return source_minus_docstrings_with_prints_cleaned, run_outputs, skipped, failed
