def test_remove_upstream_modifies_signature(backup_spec_with_functions):
# by the time we reach this test, my_tasks.raw.functions has alread been
# loaded (previous test), so we force reload to avoid wrongfully reading
# the modified source code in the raw task
from my_tasks.raw import functions
importlib.reload(functions)
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag['clean'].source.primitive
params = dag['clean'].params.to_json_serializable()
dev = CallableInteractiveDeveloper(fn, params)
nb = dev.to_nb()
# delete upstream reference
del nb.cells[-2]
dev.overwrite(nb)
source = Path('my_tasks', 'clean', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
expected = ('# adding this to make sure relative imports work '
'fine\nfrom .util import util_touch\n\n\n'
'def function(product):')
assert top_lines == expected
def test_empty_cells_at_the_end(backup_test_pkg):
dev = CallableInteractiveDeveloper(functions.simple, params={})
nb = dev.to_nb()
fmt = nbformat.versions[nbformat.current_nbformat]
nb.cells.append(fmt.new_code_cell())
path = Path(backup_test_pkg, 'functions.py')
source_old = path.read_text()
dev.overwrite(nb)
source_new = path.read_text()
assert source_old == source_new
def test_function_replace(backup_test_pkg):
path = Path(backup_test_pkg, 'functions.py')
lines = path.read_text().splitlines()
lines.insert(15, 'def impostor():\n pass\n')
path.write_text('\n'.join(lines))
dev = CallableInteractiveDeveloper(functions.simple, params={})
nb = dev.to_nb()
# anything below this must be the function's body
_, idx = find_cell_tagged(nb, 'imports-local')
fn_body = '\n'.join([c.source for c in nb.cells[idx + 1:]])
assert fn_body == ('up = upstream["some_task"]\n'
'x = 1\nPath(path).write_text(str(x))')
def test_unmodified_function(fn_name, remove_trailing_newline,
backup_test_pkg):
"""
This test makes sure the file is not modified if we don't change the
notebook because whitespace is tricky
"""
fn = getattr(functions, fn_name)
path_to_file = Path(inspect.getfile(fn))
content = path_to_file.read_text()
# make sure the trailing newline in the file is not removed accidentally,
# we need it as part of the test
assert content[-1] == '\n', 'expected a trailing newline character'
if remove_trailing_newline:
path_to_file.write_text(content[:-1])
functions_reloaded = importlib.reload(functions)
fn = getattr(functions_reloaded, fn_name)
fn_source_original = inspect.getsource(fn)
mod_source_original = path_to_file.read_text()
with CallableInteractiveDeveloper(getattr(functions_reloaded, fn_name), {
'upstream': None,
'product': None
}) as tmp_nb:
pass
functions_edited = importlib.reload(functions)
fn_source_new = inspect.getsource(getattr(functions_edited, fn_name))
mod_source_new = path_to_file.read_text()
assert fn_source_original == fn_source_new
assert mod_source_original == mod_source_new
assert not Path(tmp_nb).exists()
def test_add_upstream_modifies_signature(backup_spec_with_functions):
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag['raw'].source.primitive
params = dag['raw'].params.to_json_serializable()
dev = CallableInteractiveDeveloper(fn, params)
# add an upstream reference...
nb = dev.to_nb()
nb.cells[-1]['source'] += '\nupstream["some_task"]'
dev.overwrite(nb)
# source must be updated...
source = Path('my_tasks', 'raw', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
expected = (
'from pathlib import Path\n\n\n'
'def function(product, upstream):\n Path(str(product)).touch()')
assert expected == top_lines
# if we save again, nothing should change
dev.overwrite(nb)
source = Path('my_tasks', 'raw', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
assert expected == top_lines
def test_to_nb():
dev = CallableInteractiveDeveloper(functions.large_function, params={})
nb = dev.to_nb()
# re-construct source generated by the notebook
# this notebook has local imports anything below that should be the
# function body
_, index = find_cell_tagged(nb, 'imports-local')
source_from_nb = '\n'.join(c['source'] for c in nb.cells[index + 1:])
# get original source code
source_fn = inspect.getsource(functions.large_function)
# remove indentation from the function body
source_fn_lines = [line[4:] for line in source_fn.splitlines()]
# ignore first line from function signature
source_expected = '\n'.join(source_fn_lines[1:])
assert source_expected == source_from_nb
def test_changes_cwd(backup_test_pkg):
params = {'upstream': None, 'product': None}
source = None
with CallableInteractiveDeveloper(functions.simple, params) as tmp_nb:
nb = nbformat.read(tmp_nb, as_version=nbformat.NO_CONVERT)
cell, _ = find_cell_tagged(nb, 'debugging-settings')
source = cell.source
assert chdir_code(Path('.').resolve()) in source
def test_move_function_down(backup_test_pkg):
dev = CallableInteractiveDeveloper(functions.simple, params={})
nb = dev.to_nb()
nb.cells[-2]['source'] = 'x = 2'
# move source code down
path = Path(backup_test_pkg, 'functions.py')
source = path.read_text()
path.write_text('\n' + source)
dev.overwrite(nb)
importlib.reload(functions)
source_fn = inspect.getsource(functions.simple)
assert source_fn == ('def simple(upstream, product, path):\n '
'up = upstream["some_task"]\n '
'x = 2\n Path(path).write_text(str(x))\n')
def test_signature_line_break(backup_test_pkg):
dev = CallableInteractiveDeveloper(functions.simple, params={})
nb = dev.to_nb()
nb.cells[-2]['source'] = 'x = 2'
path = Path(backup_test_pkg, 'functions.py')
source = path.read_text()
lines = source.splitlines()
lines[16] = 'def simple(upstream, product,\npath):'
path.write_text('\n'.join(lines))
dev.overwrite(nb)
importlib.reload(functions)
source_fn = inspect.getsource(functions.simple)
assert source_fn == ('def simple(upstream, product,\npath):\n '
'up = upstream["some_task"]\n'
' x = 2\n Path(path).write_text(str(x))\n')
def test_added_imports(backup_test_pkg):
params = {'upstream': None, 'product': None}
added_imports = 'import os\nimport sys\n'
with CallableInteractiveDeveloper(functions.simple, params) as tmp_nb:
nb = nbformat.read(tmp_nb, as_version=nbformat.NO_CONVERT)
cell, _ = find_cell_tagged(nb, 'imports-top')
cell.source += f'\n{added_imports}'
nbformat.write(nb, tmp_nb)
content = Path(inspect.getfile(functions.simple)).read_text()
assert added_imports in content
def test_error_if_source_is_modified_while_editing(backup_test_pkg):
path_to_file = Path(inspect.getfile(functions.simple))
with pytest.raises(ValueError) as excinfo:
with CallableInteractiveDeveloper(functions.simple, {
'upstream': None,
'product': None
}) as nb:
path_to_file.write_text('')
assert ('Changes from the notebook were not saved back to the module'
in str(excinfo.value))
assert Path(nb).exists()
def test_develop_spec_with_local_functions(task_name,
backup_spec_with_functions):
"""
Check we can develop functions defined locally, the sample project includes
relative imports, which should work when generating the temporary notebook
"""
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag[task_name].source.primitive
params = dag[task_name].params.to_json_serializable()
if sys.platform == 'win32':
# edge case, wee need this to correctly parametrize the notebook
# when running the test on windows
params['product'] = str(params['product']).replace('\\', '\\\\')
with CallableInteractiveDeveloper(fn, params) as tmp_nb:
pm.execute_notebook(tmp_nb, tmp_nb)
def test_editing_function(fn_name, switch_indent, tmp_file, backup_test_pkg):
if switch_indent:
p = Path(functions.__file__)
p.write_text(p.read_text().replace(' ', '\t'))
functions_ = importlib.reload(functions)
else:
functions_ = functions
with CallableInteractiveDeveloper(getattr(functions_, fn_name), {
'upstream': None,
'product': None
}) as tmp_nb:
nb = nbformat.read(tmp_nb, as_version=nbformat.NO_CONVERT)
replace_first_cell(nb, 'x = 1', 'x = 2')
nbformat.write(nb, tmp_nb)
reloaded = importlib.reload(functions)
getattr(reloaded, fn_name)({'some_task': None}, None, tmp_file)
assert Path(tmp_file).read_text() == '2'
assert not Path(tmp_nb).exists()
def test_hot_reload(backup_test_pkg):
params = dict(a=1)
dev = CallableInteractiveDeveloper(functions.some_function, params=params)
# edit params
params['a'] = 2
# exporting should be aware of the change above
nb = dev.to_nb()
# edit notebook source
nb.cells[-1]['source'] = '1 + 1'
dev.overwrite(nb)
nb_new = dev.to_nb()
# check exported params match edited value
cell, _ = find_cell_tagged(nb_new, 'injected-parameters')
assert 'a = 2' in cell['source']
# should have the new content
assert nb_new.cells[-1]['source'] == '1 + 1'
def _interactive_developer(self):
"""
Creates a CallableInteractiveDeveloper instance
"""
# TODO: resolve to absolute to make relative paths work
return CallableInteractiveDeveloper(self.source.primitive, self.params)