Verdi: Warn for already existing code@computer label

aiida/cmdline/commands/cmd_code.py

@@ -64,6 +64,7 @@ def set_code_builder(ctx, param, value):
 @options_code.INPUT_PLUGIN()
 @options_code.ON_COMPUTER()
 @options_code.COMPUTER()
+@options_code.LABEL_AFTER_DUPLICATE()
 @options_code.REMOTE_ABS_PATH()
 @options_code.FOLDER()
 @options_code.REL_PATH()
@@ -104,6 +105,7 @@ def setup_code(non_interactive, **kwargs):
 @options_code.INPUT_PLUGIN(contextual_default=partial(get_default, 'input_plugin'))
 @options_code.ON_COMPUTER(contextual_default=get_on_computer)
 @options_code.COMPUTER(contextual_default=get_computer_name)
+@options_code.LABEL_AFTER_DUPLICATE(contextual_default=partial(get_default, 'label'))
 @options_code.REMOTE_ABS_PATH(contextual_default=partial(get_default, 'remote_abs_path'))
 @options_code.FOLDER(contextual_default=partial(get_default, 'code_folder'))
 @options_code.REL_PATH(contextual_default=partial(get_default, 'code_rel_path'))

aiida/cmdline/params/options/commands/code.py

@@ -13,6 +13,7 @@
 from aiida.cmdline.params import options, types
 from aiida.cmdline.params.options.interactive import InteractiveOption, TemplateInteractiveOption
 from aiida.cmdline.params.options.overridable import OverridableOption
+from aiida.cmdline.utils import echo
 
 
 def is_on_computer(ctx):
@@ -23,6 +24,88 @@ def is_not_on_computer(ctx):
     return bool(not is_on_computer(ctx))
 
 
+def label_is_none(ctx):
+    return ctx.params.get('label') is None
+
+
+def get_code_labels_of_computer(computer_node):
+    """
+    Retrieve the list of codes for the given computer in the DB.
+    """
+    from aiida.orm import Code
+    from aiida.orm.querybuilder import QueryBuilder
+    builder = QueryBuilder()
+    builder.append(entity_type='computer', filters={'id': {'==': computer_node.id}}, tag='computer')
+    builder.append(Code, project=['label'], with_computer='computer')
+    if builder.count() > 0:
+        return next(zip(*builder.all()))  # return the first entry
+
+    return []
+
+
+def get_local_code_labels():
+    """
+    Retrieve the list of codes locally stored in the DB.
+    """
+    from aiida.orm import Code
+    from aiida.orm.querybuilder import QueryBuilder
+    builder = QueryBuilder()
+    builder.append(Code, project=['label'], filters={'attributes.is_local': {'==': True}})
+    if builder.count() > 0:
+        return next(zip(*builder.all()))  # return the first entry
+
+    return []
+
+
+def _check_for_duplicate_code_label(ctx, param, value):
+    """
+    Check if the given code label and computer combination is already present
+    and ask if the user wants to enter a new label
+    """
+
+    if param.name == 'computer':
+        computer = value
+        label = ctx.params.get('label')
+    elif param.name == 'label':
+        if value is None:
+            #Even if the code/computer label combination is valid
+            #the callback will be called once for the hidden LABEL_AFTER_DUPLICATE
+            #option with a None value. In this case nothing should be done
+            return
+        computer = ctx.params.get('computer')
+        label = value
+
+    existing_codes = []
+    if computer is not None:
+        computer_label = computer.label
+        existing_codes = get_code_labels_of_computer(computer)
+    elif is_not_on_computer(ctx):
+        computer_label = 'repository'
+        existing_codes = get_local_code_labels()
+
+    if label in existing_codes:
+
+        message = f"Code '{label}@{computer_label}' already exists."
+        if param.is_non_interactive(ctx):
+            raise click.BadParameter(message)
+
+        if param.name == 'computer':
+            #First usage: Ask if the label should be changed
+            #Setting the label to None will trigger the second label prompt
+            echo.echo_warning(message)
+
+            reenter_label = click.confirm('Choose a different label:', default=True)
+            if reenter_label:
+                ctx.params['label'] = None
+        elif param.name == 'label':
+            #Second usage: If we are her the user wants to change the label
+            #so we raise click.BadParameter
+            raise click.BadParameter(message)
+    elif param.name == 'label':
+        #Overwrite the old label value
+        ctx.params['label'] = value
+
+
 ON_COMPUTER = OverridableOption(
     '--on-computer/--store-in-db',
     is_eager=False,
@@ -71,6 +154,18 @@ def is_not_on_computer(ctx):
     'computer.'
 )
 
+LABEL_AFTER_DUPLICATE = options.LABEL.clone(
+    prompt='Label',
+    cls=InteractiveOption,
+    required_fn=label_is_none,
+    prompt_fn=label_is_none,
+    callback=_check_for_duplicate_code_label,
+    hidden=True,
+    expose_value=False,
+    help="This label can be used to identify the code (using 'label@computerlabel'), as long as labels are unique per "
+    'computer.'
+)
+
 DESCRIPTION = options.DESCRIPTION.clone(
     prompt='Description',
     cls=InteractiveOption,
@@ -88,6 +183,7 @@ def is_not_on_computer(ctx):
     cls=InteractiveOption,
     required_fn=is_on_computer,
     prompt_fn=is_on_computer,
+    callback=_check_for_duplicate_code_label,
     help='Name of the computer, on which the code is installed.'
 )
 

tests/cmdline/commands/test_code.py

@@ -270,6 +270,60 @@ def test_interactive_upload(clear_database_before_test, aiida_localhost, non_int
     assert isinstance(orm.Code.get_from_string(f'{label}'), orm.Code)
 
 
+@pytest.mark.parametrize('non_interactive_editor', ('sleep 1; vim -cwq',), indirect=True)
+def test_interactive_remote_duplicate_label(
+    clear_database_before_test, aiida_local_code_factory, aiida_localhost, non_interactive_editor
+):
+    """Test interactive remote code setup with already existing lab."""
+    label = 'interactive_remote_duplicate_label'
+    aiida_local_code_factory('core.arithmetic.add', '/bin/cat', label=label)
+
+    replace_label = 'interactive_remote_replace_label'
+    user_input = '\n'.join([
+        label, 'description', 'core.arithmetic.add', 'yes', aiida_localhost.label, 'yes', replace_label,
+        '/remote/abs/path'
+    ])
+    result = CliRunner().invoke(setup_code, input=user_input)
+    assert result.exception is None
+    assert isinstance(orm.Code.get_from_string(f'{replace_label}@{aiida_localhost.label}'), orm.Code)
+
+
+def test_non_interactive_remote_duplicate_label(clear_database_before_test, aiida_local_code_factory, aiida_localhost):
+    """Test non-interactive remote code setup with already existing label."""
+    label = 'non_interactive_remote_duplicate_label'
+    aiida_local_code_factory('core.arithmetic.add', '/bin/cat', label=label)
+
+    result = CliRunner().invoke(
+        setup_code, [
+            '--non-interactive', f'--label={label}', '--description=description', '--input-plugin=core.arithmetic.add',
+            '--on-computer', f'--computer={aiida_localhost.label}', '--remote-abs-path=/remote/abs/path'
+        ]
+    )
+
+    assert result.exception is not None
+    assert f"'{label}@{aiida_localhost.label}' already exists." in result.output
+
+
+@pytest.mark.parametrize('non_interactive_editor', ('sleep 1; vim -cwq',), indirect=True)
+def test_interactive_upload_duplicate_label(
+    clear_database_before_test, aiida_local_code_factory, aiida_localhost, non_interactive_editor
+):
+    """Test interactive local code setup with already existing label."""
+    label = 'interactive_upload_duplicate_label'
+    replace_label = 'interactive_upload_replace_label'
+    dirname = os.path.dirname(__file__)
+    basename = os.path.basename(__file__)
+    user_input = '\n'.join([label, 'description', 'core.arithmetic.add', 'no', dirname, basename])
+    result = CliRunner().invoke(setup_code, input=user_input)
+    assert result.exception is None
+    assert isinstance(orm.Code.get_from_string(f'{label}'), orm.Code)
+
+    user_input = '\n'.join([label, 'description', 'core.arithmetic.add', 'no', 'yes', replace_label, dirname, basename])
+    result = CliRunner().invoke(setup_code, input=user_input)
+    assert result.exception is None
+    assert isinstance(orm.Code.get_from_string(f'{replace_label}'), orm.Code)
+
+
 @pytest.mark.parametrize('non_interactive_editor', ('sleep 1; vim -cwq',), indirect=True)
 def test_mixed(clear_database_before_test, aiida_localhost, non_interactive_editor):
     """Test mixed (interactive/from config) code setup."""
@@ -297,6 +351,35 @@ def test_code_duplicate_interactive(clear_database_before_test, aiida_local_code
     assert code.get_append_text() == duplicate.get_append_text()
 
 
+@pytest.mark.parametrize('non_interactive_editor', ('sleep 1; vim -cwq',), indirect=True)
+def test_code_duplicate_interactive_duplicate_label(
+    clear_database_before_test, aiida_local_code_factory, non_interactive_editor
+):
+    """Test code duplication interactive with already existing lable."""
+    label = 'code_duplicate_interactive_duplicate_label'
+    replace_label = 'code_duplicate_interactive_replace_label'
+    user_input = f'{label}\n\n\n\n\nyes\n{replace_label}\n\n'
+    code = aiida_local_code_factory('core.arithmetic.add', '/bin/cat', label=label)
+    result = CliRunner().invoke(code_duplicate, [str(code.pk)], input=user_input)
+    assert result.exception is None, result.exception
+
+    duplicate = orm.Code.get_from_string(replace_label)
+    assert code.description == duplicate.description
+    assert code.get_prepend_text() == duplicate.get_prepend_text()
+    assert code.get_append_text() == duplicate.get_append_text()
+
+
+def test_code_duplicate_non_interactive_duplicate_label(
+    clear_database_before_test, aiida_local_code_factory, aiida_localhost
+):
+    """Test code duplication non-interactive without chaing the label."""
+    label = 'code_duplicate_interactive_duplicate_label'
+    code = aiida_local_code_factory('core.arithmetic.add', '/bin/cat', label=label)
+    result = CliRunner().invoke(code_duplicate, ['--non-interactive', str(code.pk)])
+    assert result.exception is not None
+    assert f"'{label}@{aiida_localhost.label}' already exists." in result.output
+
+
 @pytest.mark.parametrize('non_interactive_editor', ('sleep 1; vim -cwq',), indirect=True)
 def test_code_duplicate_ignore(clear_database_before_test, aiida_local_code_factory, non_interactive_editor):
     """Providing "!" to description should lead to empty description.