def main():
start_time = datetime.datetime.now()
parser = CustomParser(description='Call an entry point '
'(pipeline.yaml or dotted path to factory)',
prog='ploomber interact')
with parser:
# this command has no static args
pass
dag, _ = parser.load_from_entry_point_arg()
try:
dag.render()
except Exception:
err = ('Your dag failed to render, but you can still inspect the '
'object to debug it.\n')
telemetry.log_api("interact_error",
dag=dag,
metadata={
'type': 'dag_render_failed',
'exception': err
})
print(err)
end_time = datetime.datetime.now()
telemetry.log_api("ploomber_interact",
total_runtime=str(end_time - start_time),
dag=dag)
# NOTE: do not use embed here, we must use start_ipython, see here:
# https://github.com/ipython/ipython/issues/8918
start_ipython(argv=[], user_ns={'dag': dag})
def test_help_does_not_display_location_if_missing_entry_point():
parser = CustomParser()
with parser:
pass
assert 'Entry point\n' in parser.format_help()
def test_add_static_arguments():
parser = CustomParser()
with parser:
parser.add_argument('--static-arg', '-s')
added = {'static_arg', 's'}
assert set(parser.static_args) & added == added
def test_help_displays_location_of_located_entry_point(tmp_directory):
Path('pipeline.yaml').touch()
parser = CustomParser()
with parser:
pass
assert 'Entry point, defaults to pipeline.yaml\n' in parser.format_help()
def test_default_loaded_from_env_var(monkeypatch):
monkeypatch.setenv('ENTRY_POINT', 'dag.yaml')
monkeypatch.setattr(sys, 'argv', ['ploomber'])
parser = CustomParser()
assert parser.DEFAULT_ENTRY_POINT == 'dag.yaml'
args = parser.parse_args()
assert args.entry_point == 'dag.yaml'
def main():
parser = CustomParser(description='Make a pipeline report')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the report, defaults to pipeline.html',
default='pipeline.html')
dag, args = _custom_command(parser)
dag.to_markup(path=args.output)
print('Report saved at:', args.output)
def main():
parser = CustomParser(description='Plot a pipeline')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the plot, defaults to pipeline.png',
default='pipeline.png')
dag, args = _custom_command(parser)
dag.plot(output=args.output)
print('Plot saved at:', args.output)
def test_shows_default_value(tmp_directory, capsys):
Path('pipeline.yaml').touch()
parser = CustomParser()
with parser:
pass
parser.print_help()
captured = capsys.readouterr()
assert 'defaults to pipeline.yaml' in captured.out
def test_entry_point_from_factory_in_environment_variable(
backup_test_pkg, monkeypatch):
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setenv('ENTRY_POINT', 'test_pkg.entry.plain_function')
parser = CustomParser()
with parser:
pass
dag, _ = parser.load_from_entry_point_arg()
assert isinstance(dag, DAG)
def test_default_loaded_from_env_var(tmp_directory, monkeypatch):
Path('pipeline.yaml').touch()
Path('dag.yaml').touch()
monkeypatch.setenv('ENTRY_POINT', 'dag.yaml')
monkeypatch.setattr(sys, 'argv', ['ploomber'])
parser = CustomParser()
assert parser.DEFAULT_ENTRY_POINT == 'dag.yaml'
args = parser.parse_args()
assert args.entry_point == 'dag.yaml'
def main():
start_time = datetime.datetime.now()
parser = CustomParser(description='Show pipeline status',
prog='ploomber status')
with parser:
# this command has no static args
pass
dag, args = parser.load_from_entry_point_arg()
print(dag.status())
end_time = datetime.datetime.now()
telemetry.log_api("ploomber_status",
total_runtime=str(end_time - start_time),
dag=dag)
def test_add_static_mutually_exclusive_group(capsys):
parser = CustomParser()
with parser:
group = parser.add_mutually_exclusive_group()
group.add_argument('--one', '-o', action='store_true')
group.add_argument('--two', '-t', action='store_true')
with pytest.raises(SystemExit):
parser.parse_args(args=['-o', '-t'])
captured = capsys.readouterr()
assert 'not allowed with argument' in captured.err
def test_custom_parser_error_if_unable_to_automatically_locate_entry_point(
capsys):
parser = CustomParser()
with parser:
pass
with pytest.raises(SystemExit) as excinfo:
parser.parse_entry_point_value()
captured = capsys.readouterr()
assert excinfo.value.code == 2
assert 'Unable to find a pipeline entry point' in captured.err
def test_shows_default_value_from_env_var(tmp_directory, monkeypatch, capsys):
monkeypatch.setenv('ENTRY_POINT', 'dag.yaml')
Path('dag.yaml').touch()
parser = CustomParser()
with parser:
pass
parser.print_help()
captured = capsys.readouterr()
assert re.search(r'defaults\s+to\s+dag.yaml\s+\(ENTRY_POINT\s+env\s+var\)',
captured.out)
def main(render_only=False):
parser = CustomParser(description='Build pipeline')
with parser:
parser.add_argument('--force',
'-f',
help='Force execution by ignoring status',
action='store_true',
default=False)
parser.add_argument('--skip-upstream',
'-su',
help='Skip building upstream dependencies. '
'Only applicable when using --partially',
action='store_true',
default=False)
parser.add_argument(
'--partially',
'-p',
help='Build a pipeline partially until certain task',
default=None)
parser.add_argument(
'--debug',
'-d',
help='Drop a debugger session if an exception happens',
action='store_true',
default=False)
dag, args = _custom_command(parser)
# when using the parallel executor from the CLI, ensure we print progress
# to stdout
if isinstance(dag.executor, Parallel):
dag.executor.print_progress = True
if render_only:
dag.render()
else:
if args.partially:
report = dag.build_partially(args.partially,
force=args.force,
debug=args.debug,
skip_upstream=args.skip_upstream)
else:
report = dag.build(force=args.force, debug=args.debug)
if report:
print(report)
return dag
def test_log(tmp_nbs, monkeypatch):
mock = Mock()
monkeypatch.setattr(
sys, 'argv',
['python', '--log', 'info', '--entry-point', 'pipeline.yaml'])
monkeypatch.setattr(parsers, 'logging', mock)
parser = CustomParser()
with parser:
pass
parser.load_from_entry_point_arg()
mock.basicConfig.assert_called_with(level='INFO')
def test_doesnt_modify_error_if_unknown_reason():
def my_fn(log=True):
pass
parser = CustomParser()
arg = Mock()
arg.options_strings = ['a', 'b']
parser.add_argument = Mock(side_effect=ArgumentError(None, 'message'))
with parser:
pass
with pytest.raises(ArgumentError) as excinfo:
_add_args_from_callable(parser, my_fn)
assert 'message' == str(excinfo.value)
def test_dagspec_initialization_from_yaml_and_env(tmp_nbs, monkeypatch):
"""
DAGSpec can be initialized with a path to a spec or a dictionary, but
they have a slightly different behavior. This ensure the cli passes
the path, instead of a dictionary
"""
mock_DAGSpec = Mock(wraps=parsers.DAGSpec)
mock_default_path_to_env = Mock(wraps=parsers.default.path_to_env)
mock_EnvDict = Mock(wraps=parsers.EnvDict)
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setattr(parsers, 'DAGSpec', mock_DAGSpec)
monkeypatch.setattr(parsers.default, 'path_to_env',
mock_default_path_to_env)
monkeypatch.setattr(parsers, 'EnvDict', mock_EnvDict)
parser = CustomParser()
with parser:
pass
dag, args = _custom_command(parser)
# ensure called using the path to the yaml spec
mock_DAGSpec.assert_called_once_with('pipeline.yaml',
env=EnvDict({'sample': False}))
# and EnvDict initialized from env.yaml
mock_EnvDict.assert_called_once_with(str(Path('env.yaml').resolve()))
def test_use_here_placeholder_when_processing_file(create_env, tmp_directory,
monkeypatch):
monkeypatch.setattr(sys, 'argv', ['ploomber'])
if create_env:
Path('env.yaml').write_text("""
key: value
""")
Path('script.py').write_text("""
# + tags=["parameters"]
upstream = None
""")
Path('pipeline.yaml').write_text("""
tasks:
- source: script.py
product: "{{here}}/out.ipynb"
""")
parser = CustomParser()
with parser:
pass
_process_file_dir_or_glob(parser)
def test_custom_parser_static_args():
parser = CustomParser()
assert set(parser.static_args) == {
'h', 'help', 'log', 'l', 'entry_point', 'e', 'log_file', 'F'
}
def test_error_if_missing_entry_point_value(monkeypatch, capsys):
monkeypatch.setattr(sys, 'argv', ['ploomber', '--entry-point'])
parser = CustomParser()
with parser:
pass
with pytest.raises(SystemExit) as excinfo:
parser.parse_entry_point_value()
captured = capsys.readouterr()
assert excinfo.value.code == 2
assert ('ploomber: error: argument --entry-point/-e: expected one argument'
in captured.err)
def main():
parser = CustomParser(description='Show pipeline status')
with parser:
# this command has no static args
pass
dag, args = _custom_command(parser)
print(dag.status())
def test_process_file_or_entry_point_param_replace(argv, expected, monkeypatch,
tmp_directory):
d = {
'meta': {
'extract_product': False,
'extract_upstream': False
},
'tasks': [{
'source': 'plot.py',
'params': {
'some_param': '{{tag}}',
},
'product': 'output/plot.ipynb',
'name': 'plot'
}]
}
Path('plot.py').write_text('# + tags=["parameters"]')
with open('pipeline.yaml', 'w') as f:
yaml.dump(d, f)
with open('env.yaml', 'w') as f:
yaml.dump({'tag': 'some_value'}, f)
monkeypatch.setattr(sys, 'argv', argv)
parser = CustomParser()
with parser:
pass
dag, args = _process_file_dir_or_glob(parser)
assert dag['plot'].params['some_param'] == expected
def main():
start_time = datetime.datetime.now()
parser = CustomParser(description='Make a pipeline report',
prog='ploomber report')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the report, defaults to pipeline.html',
default='pipeline.html')
dag, args = parser.load_from_entry_point_arg()
dag.to_markup(path=args.output)
end_time = datetime.datetime.now()
telemetry.log_api("ploomber_report",
total_runtime=str(end_time - start_time),
dag=dag,
metadata={'argv': sys.argv})
print('Report saved at:', args.output)
def test_log_file_with_factory_entry_point(backup_test_pkg, monkeypatch, opt):
mock = Mock()
monkeypatch.setattr(sys, 'argv', [
'python', '--log', 'info', opt, 'my.log', '--entry-point',
'test_pkg.entry.plain_function'
])
monkeypatch.setattr(parsers, 'logging', mock)
parser = CustomParser()
with parser:
pass
parser.load_from_entry_point_arg()
mock.basicConfig.assert_called_with(level='INFO')
mock.FileHandler.assert_called_with('my.log')
mock.getLogger().addHandler.assert_called()
def test_log_file(tmp_nbs, monkeypatch, opt):
mock = Mock()
monkeypatch.setattr(sys, 'argv', [
'python', '--log', 'info', opt, 'my.log', '--entry-point',
'pipeline.yaml'
])
monkeypatch.setattr(parsers, 'logging', mock)
parser = CustomParser()
with parser:
pass
parser.load_from_entry_point_arg()
mock.basicConfig.assert_called_with(level='INFO')
mock.FileHandler.assert_called_with('my.log')
mock.getLogger().addHandler.assert_called()
def test_cli_from_param(default, monkeypatch):
def factory(param=default):
return param
monkeypatch.setattr(sys, 'argv', ['python', '--param', 'value'])
monkeypatch.setattr(test_pkg, 'mocked_factory', factory, raising=False)
parser = CustomParser()
with parser:
pass
returned, args = parser.process_factory_dotted_path(
'test_pkg.mocked_factory')
actions = {a.dest: a for a in parser._actions}
assert actions['param'].default == default
assert args.param == 'value'
assert returned == 'value'
def main():
parser = CustomParser(description='Plot a pipeline')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the plot, defaults to pipeline.png',
default=None)
dag, args = _custom_command(parser)
if args.output is not None:
output = args.output
else:
name = extract_name(args.entry_point)
output = 'pipeline.png' if name is None else f'pipeline.{name}.png'
dag.plot(output=output)
print('Plot saved at:', output)
def test_dagspec_initialization_from_yaml(tmp_nbs_nested, monkeypatch):
"""
DAGSpec can be initialized with a path to a spec or a dictionary, but
they have a slightly different behavior. This checks that we initialize
with the path
"""
mock = Mock(wraps=parsers.DAGSpec)
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setattr(parsers, 'DAGSpec', mock)
parser = CustomParser()
with parser:
pass
dag, args = parser.load_from_entry_point_arg()
mock.assert_called_once_with('pipeline.yaml')
def test_cli_from_bool_flag(default, monkeypatch):
def factory(flag: bool = default):
pass
monkeypatch.setattr(sys, 'argv', ['python', '--flag'])
monkeypatch.setattr(test_pkg, 'mocked_factory', factory, raising=False)
parser = CustomParser()
with parser:
pass
parser.process_factory_dotted_path('test_pkg.mocked_factory')
actions = {a.dest: a for a in parser._actions}
# default should match with function's signature
assert actions['flag'].default is default
# passing the flag should flip the value
assert actions['flag'].const is not default