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