Skip to content

Commit

Permalink
👌 CLI: Computer/Code export output_file optional
Browse files Browse the repository at this point in the history
  • Loading branch information
GeigerJ2 committed Jul 9, 2024
1 parent c740b99 commit de07f8f
Show file tree
Hide file tree
Showing 14 changed files with 336 additions and 81 deletions.
33 changes: 18 additions & 15 deletions src/aiida/cmdline/commands/cmd_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
###########################################################################
"""`verdi code` command."""

import pathlib
from collections import defaultdict
from functools import partial

Expand All @@ -18,6 +19,7 @@
from aiida.cmdline.params import arguments, options, types
from aiida.cmdline.params.options.commands import code as options_code
from aiida.cmdline.utils import echo, echo_tabulate
from aiida.cmdline.utils.common import generate_validate_output_file
from aiida.cmdline.utils.decorators import with_dbenv
from aiida.common import exceptions

Expand Down Expand Up @@ -234,34 +236,35 @@ def show(code):

@verdi_code.command()
@arguments.CODE()
@arguments.OUTPUT_FILE(type=click.Path(exists=False))
@click.option(
'--sort/--no-sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False)
@options.OVERWRITE()
@options.SORT()
@with_dbenv()
def export(code, output_file, sort):
def export(code, output_file, overwrite, sort):
"""Export code to a yaml file."""

import yaml

code_data = {}

for key in code.Model.model_fields.keys():
if key == 'computer':
value = getattr(code, key).label
else:
value = getattr(code, key)
value = getattr(code, key).label if key == 'computer' else getattr(code, key)

# If the attribute is not set, for example ``with_mpi`` do not export it, because the YAML won't be valid for
# use in ``verdi code create`` since ``None`` is not a valid value on the CLI.
if value is not None:
code_data[key] = str(value)

with open(output_file, 'w', encoding='utf-8') as yfhandle:
yaml.dump(code_data, yfhandle, sort_keys=sort)
try:
output_file = generate_validate_output_file(
output_file=output_file, entity_label=code.label, overwrite=overwrite, appendix=f'@{code_data["computer"]}'
)
except (FileExistsError, IsADirectoryError) as exception:
raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception

output_file.write_text(yaml.dump(code_data, sort_keys=sort))

echo.echo_success(f'Code<{code.pk}> {code.label} exported to file `{output_file}`.')


@verdi_code.command()
Expand Down
52 changes: 30 additions & 22 deletions src/aiida/cmdline/commands/cmd_computer.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from aiida.cmdline.params import arguments, options
from aiida.cmdline.params.options.commands import computer as options_computer
from aiida.cmdline.utils import echo, echo_tabulate
from aiida.cmdline.utils.common import generate_validate_output_file
from aiida.cmdline.utils.decorators import with_dbenv
from aiida.common.exceptions import EntryPointError, ValidationError
from aiida.plugins.entry_point import get_entry_point_names
Expand Down Expand Up @@ -741,16 +742,11 @@ def computer_export():

@computer_export.command('setup')
@arguments.COMPUTER()
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path))
@click.option(
'--sort/--no-sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False)
@options.OVERWRITE()
@options.SORT()
@with_dbenv()
def computer_export_setup(computer, output_file, sort):
def computer_export_setup(computer, output_file, overwrite, sort):
"""Export computer setup to a YAML file."""
import yaml

Expand All @@ -769,6 +765,14 @@ def computer_export_setup(computer, output_file, sort):
'prepend_text': computer.get_prepend_text(),
'append_text': computer.get_append_text(),
}

try:
output_file = generate_validate_output_file(
output_file=output_file, entity_label=computer.label, overwrite=overwrite, appendix='-setup'
)
except (FileExistsError, IsADirectoryError) as exception:
raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception

try:
output_file.write_text(yaml.dump(computer_setup, sort_keys=sort), 'utf-8')
except Exception as e:
Expand All @@ -783,19 +787,14 @@ def computer_export_setup(computer, output_file, sort):

@computer_export.command('config')
@arguments.COMPUTER()
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path))
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False)
@options.USER(
help='Email address of the AiiDA user from whom to export this computer (if different from default user).'
)
@click.option(
'--sort/--no-sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@options.OVERWRITE()
@options.SORT()
@with_dbenv()
def computer_export_config(computer, output_file, user, sort):
def computer_export_config(computer, output_file, user, overwrite, sort):
"""Export computer transport configuration for a user to a YAML file."""
import yaml

Expand All @@ -804,20 +803,29 @@ def computer_export_config(computer, output_file, user, sort):
f'Computer<{computer.pk}> {computer.label} configuration cannot be exported,'
' because computer has not been configured yet.'
)
else:
try:
output_file = generate_validate_output_file(
output_file=output_file, entity_label=computer.label, overwrite=overwrite, appendix='-config'
)
except (FileExistsError, IsADirectoryError) as exception:
raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception

try:
computer_configuration = computer.get_configuration(user)
output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8')
except Exception as e:

except Exception as exception:

Check warning on line 818 in src/aiida/cmdline/commands/cmd_computer.py

View check run for this annotation

Codecov / codecov/patch

src/aiida/cmdline/commands/cmd_computer.py#L818

Added line #L818 was not covered by tests
error_traceback = traceback.format_exc()
echo.CMDLINE_LOGGER.debug(error_traceback)
if user is None:
echo.echo_critical(
f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}: {e!s}.'
f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}: {exception!s}.' # noqa: E501
)
else:
echo.echo_critical(
f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}'
f' and User<{user.pk}> {user.email}: {e!s}.'
f' and User<{user.pk}> {user.email}: {exception!s}.'
)
else:
echo.echo_success(f"Computer<{computer.pk}> {computer.label} configuration exported to file '{output_file}'.")
echo.echo_success(f'Computer<{computer.pk}> {computer.label} configuration exported to file `{output_file}`.')
1 change: 1 addition & 0 deletions src/aiida/cmdline/params/options/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@
'REPOSITORY_PATH',
'SCHEDULER',
'SILENT',
'SORT',
'TIMEOUT',
'TRAJECTORY_INDEX',
'TRANSPORT',
Expand Down
10 changes: 10 additions & 0 deletions src/aiida/cmdline/params/options/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@
'REPOSITORY_PATH',
'SCHEDULER',
'SILENT',
'SORT',
'TIMEOUT',
'TRAJECTORY_INDEX',
'TRANSPORT',
Expand Down Expand Up @@ -771,3 +772,12 @@ def set_log_level(ctx, _param, value):
show_default=True,
help='Overwrite file/directory if writing to disk.',
)

SORT = OverridableOption(
'--sort/--no-sort',
'sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
22 changes: 22 additions & 0 deletions src/aiida/cmdline/utils/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,13 @@
###########################################################################
"""Common utility functions for command line commands."""

from __future__ import annotations

import logging
import os
import sys
import textwrap
from pathlib import Path
from typing import TYPE_CHECKING

from click import style
Expand Down Expand Up @@ -481,3 +484,22 @@ def build_entries(ports):

echo.echo(tabulate(table, tablefmt='plain'))
echo.echo(style('\nExit codes that invalidate the cache are marked in bold red.\n', italic=True))


def generate_validate_output_file(
output_file: Path | None, entity_label: str, appendix: str = '', overwrite: bool = False
):
"""Generate default output filename for `Code`/`Computer` export and validate."""

if output_file is None:
output_file = Path(f'{entity_label}{appendix}.yml')

if output_file.is_dir():
raise IsADirectoryError(
f'A directory with the name `{output_file.resolve()}` already exists. Remove manually and try again.'
)

if output_file.is_file() and not overwrite:
raise FileExistsError(f'File `{output_file}` already exists, use `--overwrite` to overwrite.')

return output_file
74 changes: 65 additions & 9 deletions tests/cmdline/commands/test_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -259,8 +259,8 @@ def test_code_duplicate_ignore(run_cli_command, aiida_code_installed, non_intera


@pytest.mark.usefixtures('aiida_profile_clean')
@pytest.mark.parametrize('sort', (True, False))
def test_code_export(run_cli_command, aiida_code_installed, tmp_path, file_regression, sort):
@pytest.mark.parametrize('sort_option', ('--sort', '--no-sort'))
def test_code_export(run_cli_command, aiida_code_installed, tmp_path, file_regression, sort_option):
"""Test export the code setup to str."""
prepend_text = 'module load something\n some command'
code = aiida_code_installed(
Expand All @@ -271,14 +271,11 @@ def test_code_export(run_cli_command, aiida_code_installed, tmp_path, file_regre
)
filepath = tmp_path / 'code.yml'

options = [str(code.pk), str(filepath)]
options.append('--sort' if sort else '--no-sort')

run_cli_command(cmd_code.export, options)

options = [str(code.pk), str(filepath), sort_option]
result = run_cli_command(cmd_code.export, options)
assert str(filepath) in result.output, 'Filename should be in terminal output but was not found.'
# file regression check
with open(filepath, 'r', encoding='utf-8') as fhandle:
content = fhandle.read()
content = filepath.read_text()
file_regression.check(content, extension='.yml')

# round trip test by create code from the config file
Expand All @@ -292,6 +289,65 @@ def test_code_export(run_cli_command, aiida_code_installed, tmp_path, file_regre
assert isinstance(new_code, InstalledCode)


@pytest.mark.usefixtures('aiida_profile_clean')
def test_code_export_overwrite(run_cli_command, aiida_code_installed, tmp_path):
prepend_text = 'module load something\n some command'
code = aiida_code_installed(
default_calc_job_plugin='core.arithmetic.add',
filepath_executable='/bin/cat',
label='code',
prepend_text=prepend_text,
)
filepath = tmp_path / 'code.yml'

options = [str(code.pk), str(filepath)]

# Create directory with the same name and check that command fails
filepath.mkdir()
result = run_cli_command(cmd_code.export, options, raises=True)
assert f'A directory with the name `{filepath}` already exists' in result.output
filepath.rmdir()

# Export fails if file already exists and overwrite set to False
filepath.touch()
result = run_cli_command(cmd_code.export, options, raises=True)
assert f'File `{filepath}` already exists' in result.output

# Check that overwrite actually overwrites the exported Code config with the new data
code_echo = aiida_code_installed(
default_calc_job_plugin='core.arithmetic.add',
filepath_executable='/bin/echo',
# Need to set different label, therefore manually specify the same output filename
label='code_echo',
prepend_text=prepend_text,
)

options = [str(code_echo.pk), str(filepath), '--overwrite']
run_cli_command(cmd_code.export, options)

content = filepath.read_text()
assert '/bin/echo' in content


@pytest.mark.usefixtures('aiida_profile_clean')
@pytest.mark.usefixtures('chdir_tmp_path')
def test_code_export_default_filename(run_cli_command, aiida_code_installed):
"""Test default filename being created if no argument passed."""

prepend_text = 'module load something\n some command'
code = aiida_code_installed(
default_calc_job_plugin='core.arithmetic.add',
filepath_executable='/bin/cat',
label='code',
prepend_text=prepend_text,
)

options = [str(code.pk)]
run_cli_command(cmd_code.export, options)

assert pathlib.Path('[email protected]').is_file()


@pytest.mark.parametrize('non_interactive_editor', ('vim -cwq',), indirect=True)
def test_from_config_local_file(non_interactive_editor, run_cli_command, aiida_localhost):
"""Test setting up a code from a config file on disk."""
Expand Down
8 changes: 0 additions & 8 deletions tests/cmdline/commands/test_code/test_code_export_True_.yml

This file was deleted.

Loading

0 comments on commit de07f8f

Please sign in to comment.