Process functions: Parse docstring to set input port help attribute

aiida/engine/processes/functions.py

@@ -19,6 +19,8 @@
 import typing as t
 from typing import TYPE_CHECKING
 
+import docstring_parser
+
 from aiida.common.lang import override
 from aiida.manage import get_manager
 from aiida.orm import (
@@ -282,6 +284,7 @@ def build(func: FunctionType, node_class: t.Type['ProcessNode']) -> t.Type['Func
         :return: A Process class that represents the function
 
         """
+        # pylint: disable=too-many-statements
         if not issubclass(node_class, ProcessNode) or not issubclass(node_class, FunctionCalculationMixin):
             raise TypeError('the node_class should be a sub class of `ProcessNode` and `FunctionCalculationMixin`')
 
@@ -299,6 +302,18 @@ def build(func: FunctionType, node_class: t.Type['ProcessNode']) -> t.Type['Func
             LOGGER.warning(f'function `{func.__name__}` has invalid type hints: {exception}')
             annotations = {}
 
+        try:
+            parsed_docstring = docstring_parser.parse(func.__doc__)
+        except Exception as exception:  # pylint: disable=broad-except
+            LOGGER.warning(f'function `{func.__name__}` has a docstring that could not be parsed: {exception}')
+            param_help_string = {}
+            namespace_help_string = None
+        else:
+            param_help_string = {param.arg_name: param.description for param in parsed_docstring.params}
+            namespace_help_string = parsed_docstring.short_description if parsed_docstring.short_description else ''
+            if parsed_docstring.long_description is not None:
+                namespace_help_string += f'\n\n{parsed_docstring.long_description}'
+
         for key, parameter in signature.parameters.items():
 
             if parameter.kind in [parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY]:
@@ -323,6 +338,7 @@ def _define(cls, spec):  # pylint: disable=unused-argument
 
                 annotation = annotations.get(parameter.name)
                 valid_type = infer_valid_type_from_type_annotation(annotation) or (Data,)
+                help_string = param_help_string.get(parameter.name, None)
 
                 default = parameter.default if parameter.default is not parameter.empty else UNSPECIFIED
 
@@ -348,14 +364,22 @@ def _define(cls, spec):  # pylint: disable=unused-argument
                 else:
                     indirect_default = default
 
-                spec.input(parameter.name, valid_type=valid_type, default=indirect_default, serializer=to_aiida_type)
+                spec.input(
+                    parameter.name,
+                    valid_type=valid_type,
+                    default=indirect_default,
+                    serializer=to_aiida_type,
+                    help=help_string,
+                )
 
             # Set defaults for label and description based on function name and docstring, if not explicitly defined
             port_label = spec.inputs['metadata']['label']
 
             if not port_label.has_default():
                 port_label.default = func.__name__
 
+            spec.inputs.help = namespace_help_string
+
             # If the function supports varargs or kwargs then allow dynamic inputs, otherwise disallow
             spec.inputs.dynamic = keywords is not None or varargs
 

docs/source/howto/write_workflows.rst

@@ -199,6 +199,12 @@ Calling :meth:`~plumpy.ProcessSpec.expose_inputs` for a particular ``Process`` c
         )
         spec.output('is_even', valid_type=Bool)
 
+.. note::
+
+    The exposing functionality is not just limited to ``WorkChain`` implementations but works for all process classes, such as ``CalcJob`` plugins for example.
+    It even works for process functions (i.e., ``calcfunctions`` and ``workfunctions``) since under the hood an actual ``Process`` class is generated for them on-the-fly.
+    For process functions, the ``valid_type`` and ``help`` attributes of the exposed inputs are even preserved if they could be inferred from provided function type hints and docstrings (see :ref:`type validation<topics:processes:functions:type-validation>` and :ref:`docstring parsing<topics:processes:functions:docstring-parsing>` for details).
+
 Be aware that any inputs that already exist in the namespace will be overridden.
 To prevent this, the method accepts the ``namespace`` argument, which will cause the inputs to be copied into that namespace instead of the top-level namespace.
 This is especially useful for exposing inputs since *all* processes have the ``metadata`` input.

docs/source/topics/processes/functions.rst

@@ -116,6 +116,8 @@ The link labels for the example above will therefore be ``args_0``, ``args_1`` a
 If any of these labels were to overlap with the label of a positional or keyword argument, a ``RuntimeError`` will be raised.
 In this case, the conflicting argument name needs to be changed to something that does not overlap with the automatically generated labels for the variadic arguments.
 
+.. _topics:processes:functions:type-validation:
+
 Type validation
 ===============
 
@@ -168,6 +170,31 @@ The alternative syntax for union types ``X | Y`` `as introduced by PEP 604 <http
 If a process function has invalid type hints, they will simply be ignored and a warning message is logged: ``function 'function_name' has invalid type hints``.
 This ensures backwards compatibility in the case existing process functions had invalid type hints.
 
+.. _topics:processes:functions:docstring-parsing:
+
+Docstring parsing
+=================
+
+.. versionadded:: 2.3
+
+If a process function provides a docstring, AiiDA will attempt to parse it.
+If successful, the function argument descriptions will be set as the ``help`` attributes of the input ports of the dynamically generated process specification.
+This means the descriptions of the function arguments can be retrieved programmatically from the process specification (as returned by the ``spec`` classmethod):
+
+.. include:: include/snippets/functions/parse_docstring.py
+    :code: python
+
+This particularly useful when exposing a process function in a wrapping workchain:
+
+.. include:: include/snippets/functions/parse_docstring_expose.py
+    :code: python
+
+The user can now access the input description directly through the spec of the work chain, without having to go to the process function itself.
+For example, in an interactive shell:
+
+.. include:: include/snippets/functions/parse_docstring_expose_ipython.py
+    :code: ipython
+
 Return values
 =============
 In :numref:`fig_calculation_functions_kwargs` you can see that the engine used the label ``result`` for the link connecting the calculation function node with its output node.

docs/source/topics/processes/include/snippets/functions/parse_docstring.py

@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+from aiida.engine import calcfunction
+
+
+@calcfunction
+def add(x: int, y: int):
+    """Add two integers.
+
+    :param x: Left hand operand.
+    :param y: Right hand operand.
+    """
+    return x + y
+
+assert add.spec().inputs['a'].help == 'Left hand operand.'
+assert add.spec().inputs['b'].help == 'Right hand operand.'

docs/source/topics/processes/include/snippets/functions/parse_docstring_expose.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+from aiida.engine import WorkChain, calcfunction
+
+
+@calcfunction
+def add(x: int, y: int):
+    """Add two integers.
+
+    :param x: Left hand operand.
+    :param y: Right hand operand.
+    """
+    return x + y
+
+
+class Wrapper(WorkChain):
+    """Workchain that exposes the ``add`` calcfunction."""
+
+    @classmethod
+    def define(cls, spec):
+        super().define(spec)
+        spec.expose_inputs(add)

docs/source/topics/processes/include/snippets/functions/parse_docstring_expose_ipython.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+In [1]: builder = Wrapper.get_builder()
+
+In [2]: builder.x?
+Type:        property
+String form: <property object at 0x7f93b839a900>
+Docstring:   {
+    'name': 'x',
+    'required': 'True',
+    'valid_type': "(<class 'aiida.orm.nodes.data.int.Int'>,)",
+    'help': 'Left hand operand.',
+    'is_metadata': 'False',
+    'non_db': 'False'
+}

docs/source/topics/workflows/usage.rst

@@ -534,7 +534,8 @@ At the lowest level, each workflow should perform exactly one task.
 These workflows can then be wrapped together by a "parent" workflow to create a larger logical unit.
 
 In order to make this approach manageable, it needs to be as simple as possible to glue together multiple workflows in a larger parent workflow.
-One of the tools that AiiDA provides to simplify this is the ability to *expose* the ports of another work chain.
+One of the tools that AiiDA provides to simplify this is the ability to *expose* the ports of another process class.
+This can be another ``WorkChain`` implementation, a ``CalcJob`` or even a process function (a ``calcfunction`` or ``workfunction``).
 
 .. _topics:workflows:usage:workchains:expose_inputs_outputs:
 

environment.yml

@@ -14,6 +14,7 @@ dependencies:
 - click-spinner~=0.1.8
 - click~=8.1
 - disk-objectstore~=0.6.0
+- docstring_parser
 - get-annotations~=0.1
 - python-graphviz~=0.13
 - ipython<9,>=7

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "click-spinner~=0.1.8",
     "click~=8.1",
     "disk-objectstore~=0.6.0",
+    "docstring-parser",
     "get-annotations~=0.1;python_version<'3.10'",
     "graphviz~=0.13",
     "ipython>=7,<9",

requirements/requirements-py-3.10.txt

@@ -31,6 +31,7 @@ decorator==5.1.0
 defusedxml==0.7.1
 deprecation==2.1.0
 disk-objectstore==0.6.0
+docstring-parser==0.15.0
 docutils==0.16
 entrypoints==0.3
 Flask==2.0.3

requirements/requirements-py-3.11.txt

@@ -34,6 +34,7 @@ decorator==5.1.1
 defusedxml==0.7.1
 deprecation==2.1.0
 disk-objectstore==0.6.0
+docstring-parser==0.15.0
 docutils==0.16
 emmet-core==0.39.0
 fastjsonschema==2.16.2

requirements/requirements-py-3.8.txt

@@ -31,6 +31,7 @@ decorator==5.1.0
 defusedxml==0.7.1
 deprecation==2.1.0
 disk-objectstore==0.6.0
+docstring-parser==0.15.0
 docutils==0.16
 entrypoints==0.3
 Flask==2.0.3

requirements/requirements-py-3.9.txt

@@ -31,6 +31,7 @@ decorator==5.1.0
 defusedxml==0.7.1
 deprecation==2.1.0
 disk-objectstore==0.6.0
+docstring-parser==0.15.0
 docutils==0.16
 entrypoints==0.3
 Flask==2.0.3

tests/engine/test_process_function.py

@@ -713,3 +713,39 @@ def function_type_hinting(a: t.Union[int, float]):
     assert function_type_hinting(orm.Int(1)) == 2
     assert function_type_hinting(1.0) == 2.0
     assert function_type_hinting(orm.Float(1)) == 2.0
+
+
+def test_help_text_spec_inference():
+    """Test the parsing of docstrings to define the ``help`` message of the dynamically generated input ports."""
+
+    @calcfunction
+    def function(param_a, param_b, param_c):  # pylint: disable=unused-argument
+        """Some documentation.
+
+        :param param_a: Some description.
+        :param param_b: Fantastic docstring.
+        """
+
+    input_namespace = function.spec().inputs
+
+    assert input_namespace['param_a'].help == 'Some description.'
+    assert input_namespace['param_b'].help == 'Fantastic docstring.'
+    assert input_namespace['param_c'].help is None
+
+
+def test_help_text_spec_inference_invalid_docstring(aiida_caplog, monkeypatch):
+    """Test the parsing of docstrings does not except for invalid docstrings, but simply logs a warning."""
+    import docstring_parser
+
+    def raise_exception():
+        raise RuntimeError()
+
+    monkeypatch.setattr(docstring_parser, 'parse', lambda _: raise_exception())
+
+    @calcfunction
+    def function():
+        """Docstring."""
+
+    # Now call the spec to have it parse the docstring.
+    function.spec()  # pylint: disable=expression-not-assigned
+    assert 'function `function` has a docstring that could not be parsed' in aiida_caplog.records[0].message

tests/engine/test_work_chain.py

@@ -1640,6 +1640,32 @@ def step1(self):
 
         launch.run(Child)
 
+    def test_expose_process_function(self):
+        """Test that process functions can be exposed and the port attributes are preserved."""
+
+        @calcfunction  # type: ignore[misc]
+        def test_function(a: str, b: int):  # pylint: disable=unused-argument
+            """Some calcfunction.
+
+            :param a: A string argument.
+            :param b: An integer argument.
+            """
+
+        class ExposeProcessFunctionWorkChain(WorkChain):
+
+            @classmethod
+            def define(cls, spec):
+                super().define(spec)
+                spec.expose_inputs(test_function)
+
+        input_namespace = ExposeProcessFunctionWorkChain.spec().inputs
+        assert 'a' in input_namespace
+        assert 'b' in input_namespace
+        assert input_namespace['a'].valid_type == (orm.Str,)
+        assert input_namespace['a'].help == 'A string argument.'
+        assert input_namespace['b'].valid_type == (orm.Int,)
+        assert input_namespace['b'].help == 'An integer argument.'
+
 
 @pytest.mark.requires_rmq
 class TestWorkChainMisc:

utils/dependency_management.py

@@ -28,6 +28,7 @@
 
 SETUPTOOLS_CONDA_MAPPINGS = {
     'graphviz': 'python-graphviz',
+    'docstring-parser': 'docstring_parser',
 }
 
 CONDA_IGNORE = []