CLI: add the aiida-pseudo family cutoffs set command

aiida_pseudo/cli/__init__.py

@@ -7,6 +7,7 @@
 click_completion.init()
 
 from .root import cmd_root
+from .family import cmd_family
 from .install import cmd_install, cmd_install_family, cmd_install_sssp, cmd_install_pseudo_dojo
 from .list import cmd_list
 from .show import cmd_show

aiida_pseudo/cli/family.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+"""Commands to inspect or modify the contents of pseudo potential families."""
+import json
+
+import click
+
+from aiida.cmdline.utils import decorators, echo
+
+from ..groups.mixins import RecommendedCutoffMixin
+from .params import arguments, options
+from .root import cmd_root
+
+
+@cmd_root.group('family')
+def cmd_family():
+    """Command group to inspect or modify the contents of pseudo potential families."""
+
+
+@cmd_family.group('cutoffs')
+def cmd_family_cutoffs():
+    """Command group to inspect or modify the recommended cutoffs of pseudo potential families."""
+
+
+@cmd_family_cutoffs.command('set')
+@arguments.PSEUDO_POTENTIAL_FAMILY()
+@click.argument('cutoffs', type=click.File(mode='rb'))
+@options.STRINGENCY(required=True)
+@decorators.with_dbenv()
+def cmd_family_cutoffs_set(family, cutoffs, stringency):  # noqa: D301
+    """Set the recommended cutoffs for a pseudo potential family.
+
+    The cutoffs should be provided as a JSON file through the argument `CUTOFFS` which should have the structure:
+
+    \b
+        {
+            "Ag": {
+                "cutoff_wfc": 50.0,
+                "cutoff_rho": 200.0
+            },
+            ...
+        }
+
+    where the cutoffs are expected to be in electronvolt.
+    """
+    if not isinstance(family, RecommendedCutoffMixin):
+        raise click.BadParameter(f'family `{family}` does not support recommended cutoffs to be set.')
+
+    try:
+        data = json.load(cutoffs)
+    except ValueError as exception:
+        raise click.BadParameter(f'`{cutoffs.name}` contains invalid JSON: {exception}', param_hint='CUTOFFS')
+
+    try:
+        family.set_cutoffs({stringency: data})
+    except ValueError as exception:
+        raise click.BadParameter(f'`{cutoffs.name}` contains invalid cutoffs: {exception}', param_hint='CUTOFFS')
+
+    echo.echo_success(f'set cutoffs for `{family}` with the stringency `{stringency}`.')

aiida_pseudo/cli/params/arguments.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""Reusable arguments for CLI commands."""
+from aiida.cmdline.params.arguments import OverridableArgument
+from .types import PseudoPotentialFamilyParam
+
+__all__ = ('PSEUDO_POTENTIAL_FAMILY')
+
+PSEUDO_POTENTIAL_FAMILY = OverridableArgument(
+    'family', type=PseudoPotentialFamilyParam(sub_classes=('aiida.groups:pseudo.family',))
+)

aiida_pseudo/cli/params/options.py

@@ -8,8 +8,8 @@
 from .types import PseudoPotentialFamilyTypeParam
 
 __all__ = (
-    'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'DEFAULT_STRINGENCY', 'TRACEBACK',
-    'FAMILY_TYPE', 'ARCHIVE_FORMAT'
+    'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'STRINGENCY', 'DEFAULT_STRINGENCY',
+    'TRACEBACK', 'FAMILY_TYPE', 'ARCHIVE_FORMAT'
 )
 
 VERSION = OverridableOption(

aiida_pseudo/cli/show.py

@@ -6,12 +6,12 @@
 from aiida.cmdline.utils import decorators, echo
 
 from ..groups.mixins import RecommendedCutoffMixin
-from .params import PseudoPotentialFamilyParam, options
+from .params import arguments, options
 from .root import cmd_root
 
 
 @cmd_root.command('show')
-@click.argument('family', type=PseudoPotentialFamilyParam(sub_classes=('aiida.groups:pseudo.family',)))
+@arguments.PSEUDO_POTENTIAL_FAMILY()
 @options.STRINGENCY()
 @options_core.RAW()
 @decorators.with_dbenv()

tests/cli/test_family.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=unused-argument
+"""Tests for the command `aiida-pseudo show`."""
+import json
+import pytest
+
+from aiida_pseudo.cli.family import cmd_family_cutoffs_set
+from aiida_pseudo.groups.family import CutoffsFamily
+
+
+@pytest.mark.usefixtures('clear_db')
+def test_family_cutoffs_set(run_cli_command, get_pseudo_family, tmp_path):
+    """Test the `aiida-pseudo family cutoffs set` command."""
+    family = get_pseudo_family(cls=CutoffsFamily)
+    cutoffs = {
+        'normal': {},
+        'high': {},
+    }
+
+    for element in family.elements:
+        cutoffs['normal'][element] = {'cutoff_wfc': 1.0, 'cutoff_rho': 2.0}
+        cutoffs['high'][element] = {'cutoff_wfc': 3.0, 'cutoff_rho': 6.0}
+
+    # Set only the normal cutoffs for the family
+    family.set_cutoffs({'normal': cutoffs['normal']}, 'normal')
+
+    filepath = tmp_path / 'cutoffs.json'
+
+    # Invalid JSON
+    filepath.write_text('invalid content')
+    result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath)], raises=True)
+    assert "Error: Missing option '-s' / '--stringency'" in result.output
+
+    # Invalid cutoffs structure
+    filepath.write_text(json.dumps({'Ar': {'cutoff_rho': 300}}))
+    result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath), '-s', 'high'], raises=True)
+    assert 'Error: Invalid value for CUTOFFS:' in result.output
+
+    # Set correct stringency
+    stringency = 'high'
+    filepath.write_text(json.dumps(cutoffs['high']))
+    result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath), '-s', stringency])
+    assert 'Success: set cutoffs for' in result.output
+    assert stringency in family.get_cutoff_stringencies()
+    assert family.get_cutoffs(stringency) == cutoffs[stringency]