Add the ShellCode data plugin (#20)

The `ShellCode` data plugin is a subclass of the `InstalledCode` plugin
from `aiida-core`. The new plugin is used by the `launch_shell_job`
utility function whenever a code has to be created on the fly. This will
serve to be able to query for these kinds of codes and to distinguish
them from codes that are typically setup manually.

pyproject.toml

@@ -47,6 +47,9 @@ tests = [
 [project.entry-points.'aiida.calculations']
 'core.shell' = 'aiida_shell.calculations.shell:ShellJob'
 
+[project.entry-points.'aiida.data']
+'core.code.installed.shell' = 'aiida_shell.data.code:ShellCode'
+
 [project.entry-points.'aiida.parsers']
 'core.shell' = 'aiida_shell.parsers.shell:ShellParser'
 
@@ -109,6 +112,7 @@ disable = [
     'duplicate-code',
     'import-outside-toplevel',
     'inconsistent-return-statements',
+    'too-many-ancestors',
 ]
 
 [tool.yapf]

src/aiida_shell/__init__.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """AiiDA plugin that makes running shell commands easy."""
 from .calculations import ShellJob
+from .data import ShellCode
 from .engine import launch_shell_job
 from .parsers import ShellParser
 

src/aiida_shell/calculations/shell.py

@@ -84,6 +84,9 @@ def validate_nodes(cls, value: t.Mapping[str, Data], _) -> str | None:
     @classmethod
     def validate_outputs(cls, value: List, _) -> str | None:
         """Validate the ``outputs`` input."""
+        if not value:
+            return None
+
         for reserved in [cls.FILENAME_STATUS, cls.FILENAME_STDERR, cls.FILENAME_STDOUT]:
             if reserved in value:
                 return f'`{reserved}` is a reserved output filename and cannot be used in `outputs`.'

src/aiida_shell/data/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+"""Module for :mod:`aiida_shell.data`."""
+from .code import ShellCode
+
+__all__ = ('ShellCode',)

src/aiida_shell/data/code.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+"""Code that represents a shell command."""
+from __future__ import annotations
+
+from aiida.orm import InstalledCode
+
+__all__ = ('ShellCode',)
+
+
+class ShellCode(InstalledCode):
+    """Code that represents a shell command.
+
+    This code type is automatically generated by the :func:`~aiida_shell.engine.launch.launch_shell_job` function and is
+    a subclass of :class:`~aiida.orm.nodes.data.code.installed.InstalledCode`. It can therefore be used for any other
+    calculation job as well.
+    """
+
+    def __init__(self, *args, default_calc_job_plugin: str = 'core.shell', **kwargs) -> None:
+        """Construct a new instance."""
+        self.validate_default_calc_job_plugin(default_calc_job_plugin)
+        super().__init__(*args, default_calc_job_plugin=default_calc_job_plugin, **kwargs)
+
+    @staticmethod
+    def validate_default_calc_job_plugin(default_calc_job_plugin: str) -> None:
+        """Validate the default calculation job plugin.
+
+        The ``ShellCode`` should only be used with the ``core.shell`` calculation job entry point.
+
+        :raises ValueError: If ``default_calc_job_plugin`` is not ``core.shell``.
+        """
+        if default_calc_job_plugin != 'core.shell':
+            raise ValueError(f'`default_calc_job_plugin` has to be `core.shell`, but got: {default_calc_job_plugin}')
+
+    @classmethod
+    def _get_cli_options(cls) -> dict:
+        """Return the CLI options that would allow to create an instance of this class."""
+        options = super()._get_cli_options()
+        options['default_calc_job_plugin']['default'] = 'core.shell'
+        return options

src/aiida_shell/engine/launchers/shell_job.py

@@ -9,16 +9,16 @@
 
 from aiida.common import exceptions
 from aiida.engine import launch
-from aiida.orm import Code, Computer, Data, ProcessNode, SinglefileData, load_code, load_computer
+from aiida.orm import AbstractCode, Computer, Data, ProcessNode, SinglefileData, load_code, load_computer
 
-from aiida_shell.calculations.shell import ShellJob
+from aiida_shell import ShellCode, ShellJob
 
 __all__ = ('launch_shell_job',)
 
 LOGGER = logging.getLogger('aiida_shell')
 
 
-def launch_shell_job(  # pylint: disable=too-many-arguments,too-many-locals
+def launch_shell_job(  # pylint: disable=too-many-arguments
     command: str,
     nodes: dict[str, Data] | None = None,
     filenames: dict[str, str] | None = None,
@@ -40,7 +40,36 @@ def launch_shell_job(  # pylint: disable=too-many-arguments,too-many-locals
     :raises ValueError: If the absolute path of the command on the computer could not be determined.
     :returns: The tuple of results dictionary and ``ProcessNode``, or just the ``ProcessNode`` if ``submit=True``.
     """
-    computer = prepare_computer((metadata or {}).get('options', {}).pop('computer', None))
+    computer = (metadata or {}).get('options', {}).pop('computer', None)
+    code = prepare_code(command, computer)
+
+    inputs = {
+        'code': code,
+        'nodes': convert_nodes_single_file_data(nodes or {}),
+        'filenames': filenames,
+        'arguments': arguments,
+        'outputs': outputs,
+        'metadata': metadata or {},
+    }
+
+    if submit:
+        return launch.submit(ShellJob, **inputs)
+
+    results, node = launch.run_get_node(ShellJob, **inputs)
+
+    return {label: node for label, node in results.items() if isinstance(node, SinglefileData)}, node
+
+
+def prepare_code(command: str, computer: Computer | None = None) -> AbstractCode:
+    """Prepare a code for the given command and computer.
+
+    This will automatically prepare the computer
+
+    :param command: The command that the code should represent. Can be the relative executable name or absolute path.
+    :param computer: The computer on which the command should be run. If not defined the localhost will be used.
+    :return: A :class:`aiida.orm.nodes.code.abstract.AbstractCode` instance.
+    """
+    computer = prepare_computer(computer)
 
     with computer.get_transport() as transport:
         status, stdout, stderr = transport.exec_command_wait(f'which {command}')
@@ -52,30 +81,14 @@ def launch_shell_job(  # pylint: disable=too-many-arguments,too-many-locals
     code_label = f'{command}@{computer.label}'
 
     try:
-        code = load_code(code_label)
+        code: AbstractCode = load_code(code_label)
     except exceptions.NotExistent:
         LOGGER.info('No code exists yet for `%s`, creating it now.', code_label)
-        code = Code(  # type: ignore[assignment]
-            label=command,
-            remote_computer_exec=(computer, executable),
-            input_plugin_name='core.shell'
+        code = ShellCode(  # type: ignore[assignment]
+            label=command, computer=computer, filepath_executable=executable, default_calc_job_plugin='core.shell'
         ).store()
 
-    inputs = {
-        'code': code,
-        'nodes': convert_nodes_single_file_data(nodes or {}),
-        'filenames': filenames or {},
-        'arguments': arguments or [],
-        'outputs': outputs or [],
-        'metadata': metadata or {},
-    }
-
-    if submit:
-        return launch.submit(ShellJob, **inputs)
-
-    results, node = launch.run_get_node(ShellJob, **inputs)
-
-    return {label: node for label, node in results.items() if isinstance(node, SinglefileData)}, node
+    return code
 
 
 def prepare_computer(computer: Computer | None = None) -> Computer:

tests/conftest.py

@@ -11,10 +11,12 @@
 from aiida.common.links import LinkType
 from aiida.engine.utils import instantiate_process
 from aiida.manage.manager import get_manager
-from aiida.orm import CalcJobNode, Code, Computer, FolderData
+from aiida.orm import CalcJobNode, Computer, FolderData
 from aiida.plugins import CalculationFactory, ParserFactory
 import pytest
 
+from aiida_shell import ShellCode
+
 pytest_plugins = ['aiida.manage.tests.pytest_fixtures']  # pylint: disable=invalid-name
 
 
@@ -135,10 +137,10 @@ def factory(label='localhost', hostname='localhost', scheduler_type='core.direct
 
 @pytest.fixture
 def generate_code(generate_computer):
-    """Return a :class:`aiida.orm.Code` instance, either already existing or created."""
+    """Return a :class:`aiida_shell.data.code.ShellCode` instance, either already existing or created."""
 
     def factory(command='/bin/true', computer_label='localhost', label=None, entry_point_name='core.shell'):
-        """Return a :class:`aiida.orm.Code` instance, either already existing or created."""
+        """Return a :class:`aiida_shell.data.code.ShellCode` instance, either already existing or created."""
         label = label or str(uuid.uuid4())
         computer = generate_computer(computer_label)
 
@@ -151,10 +153,14 @@ def factory(command='/bin/true', computer_label='localhost', label=None, entry_p
 
         try:
             filters = {'label': label, 'attributes.input_plugin_name': entry_point_name}
-            return Code.collection.get(**filters)
+            return ShellCode.collection.get(**filters)
         except exceptions.NotExistent:
-            code = Code(label=label, input_plugin_name=entry_point_name, remote_computer_exec=[computer, executable])
-            return code.store()
+            return ShellCode(
+                label=label,
+                computer=computer,
+                filepath_executable=executable,
+                default_calc_job_plugin=entry_point_name
+            ).store()
 
     return factory
 

tests/data/test_code.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+"""Tests for :mod:`aiida_shell.data.code`."""
+import pytest
+
+from aiida_shell.data.code import ShellCode
+
+
+def test_constructor(generate_computer):
+    """Test initializing an instance."""
+    code = ShellCode(
+        label='bash',
+        computer=generate_computer(),
+        filepath_executable='/bin/bash',
+        default_calc_job_plugin='core.shell',
+    )
+    assert isinstance(code, ShellCode)
+
+
+def test_constructor_invalid(generate_computer):
+    """Test the constructor raises if ``default_calc_job_plugin`` is not ``core.shell``."""
+    with pytest.raises(ValueError, match=r'`default_calc_job_plugin` has to be `core.shell`, but got: .*'):
+        ShellCode(
+            label='bash',
+            computer=generate_computer(),
+            filepath_executable='/bin/bash',
+            default_calc_job_plugin='core.arithmetic.add',
+        )
+
+
+@pytest.mark.parametrize(('value', 'exception'), (
+    ('core.shell', None),
+    ('core.arithmetic.add', r'`default_calc_job_plugin` has to be `core.shell`, but got: .*'),
+))
+def test_validate_default_calc_job_plugin(value, exception):
+    """Test the constructor raises if ``default_calc_job_plugin`` is not ``core.shell``."""
+    if exception:
+        with pytest.raises(ValueError, match=exception):
+            ShellCode.validate_default_calc_job_plugin(value)
+    else:
+        assert ShellCode.validate_default_calc_job_plugin(value) is None