# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
"""Common cli utilities for transport plugins."""
from functools import partial
import inspect
import click
from aiida.cmdline.params import arguments, options
from aiida.cmdline.params.options.interactive import InteractiveOption
from aiida.cmdline.utils import echo
from aiida.cmdline.utils.decorators import with_dbenv
from aiida.common.exceptions import NotExistent
from aiida.manage import get_manager
TRANSPORT_PARAMS = []
# pylint: disable=unused-argument
[docs]def match_comp_transport(ctx, param, computer, transport_type):
"""Check the computer argument against the transport type."""
if computer.transport_type != transport_type:
echo.echo_critical(
f'Computer {computer.label} has transport of type "{computer.transport_type}", not {transport_type}!'
)
return computer
[docs]@with_dbenv()
def configure_computer_main(computer, user, **kwargs):
"""Configure a computer via the CLI."""
from aiida import orm
user = user or orm.User.objects.get_default()
echo.echo_report(f'Configuring computer {computer.label} for user {user.email}.')
if user.email != get_manager().get_profile().default_user_email:
echo.echo_report('Configuring different user, defaults may not be appropriate.')
computer.configure(user=user, **kwargs)
echo.echo_success(f'{computer.label} successfully configured for {user.email}')
[docs]def common_params(command_func):
"""Decorate a command function with common click parameters for all transport plugins."""
for param in TRANSPORT_PARAMS.copy().reverse():
command_func = param(command_func)
return command_func
[docs]def transport_option_default(name, computer):
"""Determine the default value for an auth_param key."""
transport_cls = computer.get_transport_class()
suggester_name = f'_get_{name}_suggestion_string'
members = dict(inspect.getmembers(transport_cls))
suggester = members.get(suggester_name, None)
default = None
if suggester:
default = suggester(computer)
else:
default = transport_cls.auth_options[name].get('default')
return default
[docs]def interactive_default(key, also_non_interactive=False):
"""Create a contextual_default value callback for an auth_param key.
:param key: the name of the option.
:param also_non_interactive: indicates whether this option should provide a default also in non-interactive mode. If
False, the option will raise `MissingParameter` if no explicit value is specified when the command is called in
non-interactive mode.
"""
@with_dbenv()
def get_default(ctx):
"""Determine the default value from the context."""
from aiida import orm
if not also_non_interactive and ctx.params['non_interactive']:
raise click.MissingParameter()
user = ctx.params.get('user', None) or orm.User.objects.get_default()
computer = ctx.params.get('computer', None)
if computer is None:
return None
try:
authinfo = orm.AuthInfo.objects.get(dbcomputer_id=computer.id, aiidauser_id=user.id)
except NotExistent:
authinfo = orm.AuthInfo(computer=computer, user=user)
auth_params = authinfo.get_auth_params()
suggestion = auth_params.get(key)
suggestion = suggestion or transport_option_default(key, computer)
return suggestion
return get_default
[docs]def create_option(name, spec):
"""Create a click option from a name and partial specs as used in transport auth_options."""
from copy import deepcopy
spec = deepcopy(spec)
name_dashed = name.replace('_', '-')
option_name = f'--{name_dashed}'
existing_option = spec.pop('option', None)
if spec.pop('switch', False):
option_name = '--{name}/--no-{name}'.format(name=name_dashed)
kwargs = {'cls': InteractiveOption, 'show_default': True}
non_interactive_default = spec.pop('non_interactive_default', False)
kwargs['contextual_default'] = interactive_default(name, also_non_interactive=non_interactive_default)
kwargs.update(spec)
if existing_option:
return existing_option(**kwargs)
return click.option(option_name, **kwargs)
[docs]def list_transport_options(transport_type):
from aiida.plugins import TransportFactory
options_list = [create_option(*item) for item in TransportFactory(transport_type).auth_options.items()]
return options_list
[docs]def transport_options(transport_type):
"""Decorate a command with all options for a computer configure subcommand for transport_type."""
def apply_options(func):
"""Decorate the command functionn with the appropriate options for the transport type."""
options_list = list_transport_options(transport_type)
options_list.reverse()
func = arguments.COMPUTER(callback=partial(match_comp_transport, transport_type=transport_type))(func)
func = options.NON_INTERACTIVE()(func)
for option in options_list:
func = option(func)
func = options.USER()(func)
func = options.CONFIG_FILE()(func)
return func
return apply_options