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 main():
parser = CustomParser(description='Manage scripts and notebooks',
prog='ploomber nb')
with parser:
# The next options do not require a valid entry point
# opening .py files as notebooks in JupyterLab with a single click
single_click = parser.add_mutually_exclusive_group()
single_click.add_argument(
'--single-click',
'-S',
action='store_true',
help=('Override JupyterLab defaults to open '
'scripts as notebook with a single click'))
single_click.add_argument(
'--single-click-disable',
'-d',
action='store_true',
help=('Disables opening scripts as notebook with a single '
'click in JupyterLab'))
# install/uninstall hook
hook = parser.add_mutually_exclusive_group()
hook.add_argument('--install-hook',
'-I',
action='store_true',
help='Install git pre-commit hook')
hook.add_argument('--uninstall-hook',
'-u',
action='store_true',
help='Uninstall git pre-commit hook')
# The next options require a valid entry point
# inject/remove cell
cell = parser.add_mutually_exclusive_group()
cell.add_argument('--inject',
'-i',
action='store_true',
help='Inject cell to all script/notebook tasks')
cell.add_argument(
'--remove',
'-r',
action='store_true',
help='Remove injected cell in all script/notebook tasks')
# re-format
parser.add_argument('--format',
'-f',
help='Re-format all script/notebook tasks')
# pair scripts and nbs
parser.add_argument('--pair',
'-p',
help='Pair scripts with ipynb files')
# sync scripts and nbs
parser.add_argument('--sync',
'-s',
action='store_true',
help='Sync scripts with ipynb files')
loading_error = None
# commands that need an entry point to work
needs_entry_point = {'format', 'inject', 'remove', 'sync', 'pair'}
args_ = parser.parse_args()
if any(getattr(args_, arg) for arg in needs_entry_point):
try:
dag, args = parser.load_from_entry_point_arg()
except Exception as e:
loading_error = e
else:
dag.render(show_progress=False)
if loading_error:
err = ('Could not run nb command: the DAG ' 'failed to load')
telemetry.log_api("nb_error",
metadata={
'type': 'dag_load_failed',
'exception': err + f' {loading_error}',
'argv': sys.argv
})
raise RuntimeError(err) from loading_error
else:
dag = None
args = args_
# options that do not need a DAG
if args.single_click:
_py_with_single_click_enable()
if args.single_click_disable:
_py_with_single_click_disable()
if args.install_hook:
if not Path('.git').is_dir():
err = ('Expected a .git/ directory in the current working '
'directory. Run this from the repository root directory.')
telemetry.log_api("nb_error",
metadata={
'type': 'no_git_config',
'exception': err,
'argv': sys.argv
})
raise NotADirectoryError(err)
parent = Path('.git', 'hooks')
parent.mkdir(exist_ok=True)
# pre-commit: remove injected cells
_install_hook(parent / 'pre-commit', pre_commit_hook, args.entry_point)
click.echo('Successfully installed pre-commit git hook')
# post-commit: inject cells
_install_hook(parent / 'post-commit', post_commit_hook,
args.entry_point)
click.echo('Successfully installed post-commit git hook')
if args.uninstall_hook:
_delete_hook(Path('.git', 'hooks', 'pre-commit'))
_delete_hook(Path('.git', 'hooks', 'post-commit'))
# options that need a valid DAG
if args.format:
new_paths = [
str(p) for p in _call_in_source(
dag,
'format',
'Formatted notebooks',
dict(fmt=args.format),
) if p is not None
]
if len(new_paths):
click.echo('Extension changed for the following '
f'tasks: {", ".join(new_paths)}. Update your '
'pipeline declaration.')
if args.inject:
_call_in_source(
dag,
'save_injected_cell',
'Injected celll',
dict(),
)
click.secho(
'Finished cell injection. Re-run this command if your '
'pipeline.yaml changes.',
fg='green')
if args.remove:
_call_in_source(
dag,
'remove_injected_cell',
'Removed injected cell',
dict(),
)
if args.sync:
# maybe its more efficient to pass all notebook paths at once?
_call_in_source(dag, 'sync', 'Synced notebooks')
# can pair give trouble if we're reformatting?
if args.pair:
_call_in_source(
dag,
'pair',
'Paired notebooks',
dict(base_path=args.pair),
)
click.echo(f'Finished pairing notebooks. Tip: add {args.pair!r} to '
'your .gitignore to keep your repository clean')
telemetry.log_api("ploomber_nb", dag=dag, metadata={'argv': sys.argv})