Process functions: Infer argument valid_type from type hints

.pre-commit-config.yaml

@@ -110,7 +110,6 @@ repos:
                 aiida/engine/processes/calcjobs/monitors.py|
                 aiida/engine/processes/calcjobs/tasks.py|
                 aiida/engine/processes/control.py|
-                aiida/engine/processes/functions.py|
                 aiida/engine/processes/ports.py|
                 aiida/manage/configuration/__init__.py|
                 aiida/manage/configuration/config.py|

aiida/engine/processes/functions.py

@@ -15,23 +15,49 @@
 import inspect
 import logging
 import signal
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Sequence, Tuple, Type, TypeVar
+import types
+import typing as t
+from typing import TYPE_CHECKING
 
 from aiida.common.lang import override
 from aiida.manage import get_manager
-from aiida.orm import CalcFunctionNode, Data, ProcessNode, WorkFunctionNode, to_aiida_type
+from aiida.orm import (
+    Bool,
+    CalcFunctionNode,
+    Data,
+    Dict,
+    Float,
+    Int,
+    List,
+    ProcessNode,
+    Str,
+    WorkFunctionNode,
+    to_aiida_type,
+)
 from aiida.orm.utils.mixins import FunctionCalculationMixin
 
 from .process import Process
 
+try:
+    UnionType = types.UnionType  # type: ignore[attr-defined]
+except AttributeError:
+    # This type is not available for Python 3.9 and older
+    UnionType = None  # pylint: disable=invalid-name
+
+try:
+    get_annotations = inspect.get_annotations  # type: ignore[attr-defined]
+except AttributeError:
+    # This is the backport for Python 3.9 and older
+    from get_annotations import get_annotations  # type: ignore[no-redef]
+
 if TYPE_CHECKING:
     from .exit_code import ExitCode
 
 __all__ = ('calcfunction', 'workfunction', 'FunctionProcess')
 
 LOGGER = logging.getLogger(__name__)
 
-FunctionType = TypeVar('FunctionType', bound=Callable[..., Any])
+FunctionType = t.TypeVar('FunctionType', bound=t.Callable[..., t.Any])
 
 
 def calcfunction(function: FunctionType) -> FunctionType:
@@ -88,14 +114,14 @@ def workfunction(function: FunctionType) -> FunctionType:
     return process_function(node_class=WorkFunctionNode)(function)
 
 
-def process_function(node_class: Type['ProcessNode']) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+def process_function(node_class: t.Type['ProcessNode']) -> t.Callable[[FunctionType], FunctionType]:
     """
     The base function decorator to create a FunctionProcess out of a normal python function.
 
     :param node_class: the ORM class to be used as the Node record for the FunctionProcess
     """
 
-    def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
+    def decorator(function: FunctionType) -> FunctionType:
         """
         Turn the decorated function into a FunctionProcess.
 
@@ -104,7 +130,7 @@ def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
         """
         process_class = FunctionProcess.build(function, node_class=node_class)
 
-        def run_get_node(*args, **kwargs) -> Tuple[Optional[Dict[str, Any]], 'ProcessNode']:
+        def run_get_node(*args, **kwargs) -> tuple[dict[str, t.Any] | None, 'ProcessNode']:
             """
             Run the FunctionProcess with the supplied inputs in a local runner.
 
@@ -159,7 +185,7 @@ def kill_process(_num, _frame):
 
             return result, process.node
 
-        def run_get_pk(*args, **kwargs) -> Tuple[Optional[Dict[str, Any]], int]:
+        def run_get_pk(*args, **kwargs) -> tuple[dict[str, t.Any] | None, int]:
             """Recreate the `run_get_pk` utility launcher.
 
             :param args: input arguments to construct the FunctionProcess
@@ -185,15 +211,52 @@ def decorated_function(*args, **kwargs):
         decorated_function.recreate_from = process_class.recreate_from  # type: ignore[attr-defined]
         decorated_function.spec = process_class.spec  # type: ignore[attr-defined]
 
-        return decorated_function
+        return decorated_function  # type: ignore[return-value]
 
     return decorator
 
 
+def infer_valid_type_from_type_annotation(annotation: t.Any) -> tuple[t.Any, ...]:
+    """Infer the value for the ``valid_type`` of an input port from the given function argument annotation.
+
+    :param annotation: The annotation of a function argument as returned by ``inspect.get_annotation``.
+    :returns: A tuple of valid types. If no valid types were defined or they could not be successfully parsed, an empty
+        tuple is returned.
+    """
+
+    def get_type_from_annotation(annotation):
+        valid_type_map = {
+            bool: Bool,
+            dict: Dict,
+            t.Dict: Dict,
+            float: Float,
+            int: Int,
+            list: List,
+            t.List: List,
+            str: Str,
+        }
+
+        if inspect.isclass(annotation) and issubclass(annotation, Data):
+            return annotation
+
+        return valid_type_map.get(annotation)
+
+    inferred_valid_type: tuple[t.Any, ...] = ()
+
+    if inspect.isclass(annotation):
+        inferred_valid_type = (get_type_from_annotation(annotation),)
+    elif t.get_origin(annotation) is t.Union or t.get_origin(annotation) is UnionType:
+        inferred_valid_type = tuple(get_type_from_annotation(valid_type) for valid_type in t.get_args(annotation))
+    elif t.get_origin(annotation) is t.Optional:
+        inferred_valid_type = (t.get_args(annotation),)
+
+    return tuple(valid_type for valid_type in inferred_valid_type if valid_type is not None)
+
+
 class FunctionProcess(Process):
     """Function process class used for turning functions into a Process"""
 
-    _func_args: Sequence[str] = ()
+    _func_args: t.Sequence[str] = ()
     _varargs: str | None = None
 
     @staticmethod
@@ -205,7 +268,7 @@ def _func(*_args, **_kwargs) -> dict:
         return {}
 
     @staticmethod
-    def build(func: Callable[..., Any], node_class: Type['ProcessNode']) -> Type['FunctionProcess']:
+    def build(func: FunctionType, node_class: t.Type['ProcessNode']) -> t.Type['FunctionProcess']:
         """
         Build a Process from the given function.
 
@@ -222,48 +285,70 @@ def build(func: Callable[..., Any], node_class: Type['ProcessNode']) -> Type['Fu
         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`')
 
-        args, varargs, keywords, defaults, _, _, _ = inspect.getfullargspec(func)
-        nargs = len(args)
-        ndefaults = len(defaults) if defaults else 0
-        first_default_pos = nargs - ndefaults
+        signature = inspect.signature(func)
+
+        args: list[str] = []
+        varargs: str | None = None
+        keywords: str | None = None
+
+        try:
+            annotations = get_annotations(func, eval_str=True)
+        except Exception as exception:  # pylint: disable=broad-except
+            # Since we are running with ``eval_str=True`` to unstringize the annotations, the call can except if the
+            # annotations are incorrect. In this case we simply want to log a warning and continue with type inference.
+            LOGGER.warning(f'function `{func.__name__}` has invalid type hints: {exception}')
+            annotations = {}
+
+        for key, parameter in signature.parameters.items():
+
+            if parameter.kind in [parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY]:
+                args.append(key)
+
+            if parameter.kind is parameter.VAR_POSITIONAL:
+                varargs = key
+
+            if parameter.kind is parameter.VAR_KEYWORD:
+                varargs = key
 
         def _define(cls, spec):  # pylint: disable=unused-argument
             """Define the spec dynamically"""
             from plumpy.ports import UNSPECIFIED
 
             super().define(spec)
 
-            for i, arg in enumerate(args):
+            for parameter in signature.parameters.values():
+
+                if parameter.kind in [parameter.VAR_POSITIONAL, parameter.VAR_KEYWORD]:
+                    continue
 
-                default = UNSPECIFIED
+                annotation = annotations.get(parameter.name)
+                valid_type = infer_valid_type_from_type_annotation(annotation) or (Data,)
 
-                if defaults and i >= first_default_pos:
-                    default = defaults[i - first_default_pos]
+                default = parameter.default if parameter.default is not parameter.empty else UNSPECIFIED
 
                 # If the keyword was already specified, simply override the default
-                if spec.has_input(arg):
-                    spec.inputs[arg].default = default
+                if spec.has_input(parameter.name):
+                    spec.inputs[parameter.name].default = default
+                    continue
+
+                # If the default is ``None`` make sure that the port also accepts a ``NoneType``. Note that we cannot
+                # use ``None`` because the validation will call ``isinstance`` which does not work when passing ``None``
+                # but it does work with ``NoneType`` which is returned by calling ``type(None)``.
+                if default is None:
+                    valid_type += (type(None),)
+
+                # If a default is defined and it is not a ``Data`` instance it should be serialized, but this should be
+                # done lazily using a lambda, just as any port defaults should not define node instances directly as is
+                # also checked by the ``spec.input`` call.
+                if (
+                    default is not None and default != UNSPECIFIED and not isinstance(default, Data) and
+                    not callable(default)
+                ):
+                    indirect_default = lambda value=default: to_aiida_type(value)
                 else:
-                    # If the default is `None` make sure that the port also accepts a `NoneType`
-                    # Note that we cannot use `None` because the validation will call `isinstance` which does not work
-                    # when passing `None`, but it does work with `NoneType` which is returned by calling `type(None)`
-                    if default is None:
-                        valid_type = (Data, type(None))
-                    else:
-                        valid_type = (Data,)
-
-                    # If a default is defined and it is not a ``Data`` instance it should be serialized, but this should
-                    # be done lazily using a lambda, just as any port defaults should not define node instances directly
-                    # as is also checked by the ``spec.input`` call.
-                    if (
-                        default is not None and default != UNSPECIFIED and not isinstance(default, Data) and
-                        not callable(default)
-                    ):
-                        indirect_default = lambda value=default: to_aiida_type(value)
-                    else:
-                        indirect_default = default  # type: ignore[assignment]
-
-                    spec.input(arg, valid_type=valid_type, default=indirect_default, serializer=to_aiida_type)
+                    indirect_default = default
+
+                spec.input(parameter.name, valid_type=valid_type, default=indirect_default, serializer=to_aiida_type)
 
             # Set defaults for label and description based on function name and docstring, if not explicitly defined
             port_label = spec.inputs['metadata']['label']
@@ -293,7 +378,7 @@ def _define(cls, spec):  # pylint: disable=unused-argument
         )
 
     @classmethod
-    def validate_inputs(cls, *args: Any, **kwargs: Any) -> None:  # pylint: disable=unused-argument
+    def validate_inputs(cls, *args: t.Any, **kwargs: t.Any) -> None:  # pylint: disable=unused-argument
         """
         Validate the positional and keyword arguments passed in the function call.
 
@@ -314,7 +399,7 @@ def validate_inputs(cls, *args: Any, **kwargs: Any) -> None:  # pylint: disable=
             raise TypeError(f'{name}() takes {nparameters} positional arguments but {nargs} were given')
 
     @classmethod
-    def create_inputs(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+    def create_inputs(cls, *args: t.Any, **kwargs: t.Any) -> dict[str, t.Any]:
         """Create the input args for the FunctionProcess."""
         cls.validate_inputs(*args, **kwargs)
 
@@ -326,7 +411,7 @@ def create_inputs(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
         return ins
 
     @classmethod
-    def args_to_dict(cls, *args: Any) -> Dict[str, Any]:
+    def args_to_dict(cls, *args: t.Any) -> dict[str, t.Any]:
         """
         Create an input dictionary (of form label -> value) from supplied args.
 
@@ -375,7 +460,7 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(enable_persistence=False, *args, **kwargs)  # type: ignore
 
     @property
-    def process_class(self) -> Callable[..., Any]:
+    def process_class(self) -> t.Callable[..., t.Any]:
         """
         Return the class that represents this Process, for the FunctionProcess this is the function itself.
 
@@ -388,7 +473,7 @@ class that really represents what was being executed.
         """
         return self._func
 
-    def execute(self) -> Optional[Dict[str, Any]]:
+    def execute(self) -> dict[str, t.Any] | None:
         """Execute the process."""
         result = super().execute()
 
@@ -405,7 +490,7 @@ def _setup_db_record(self) -> None:
         self.node.store_source_info(self._func)
 
     @override
-    def run(self) -> Optional['ExitCode']:
+    def run(self) -> 'ExitCode' | None:
         """Run the process."""
         from .exit_code import ExitCode
 
@@ -414,7 +499,7 @@ def run(self) -> Optional['ExitCode']:
         # been overridden by the engine to `Running` so we cannot check that, but if the `exit_status` is anything other
         # than `None`, it should mean this node was taken from the cache, so the process should not be rerun.
         if self.node.exit_status is not None:
-            return self.node.exit_status
+            return ExitCode(self.node.exit_status, self.node.exit_message)
 
         # Split the inputs into positional and keyword arguments
         args = [None] * len(self._func_args)

docs/source/topics/calculations/concepts.rst

@@ -55,6 +55,8 @@ To solve this, one only has to wrap them in the :py:class:`~aiida.orm.nodes.data
 The only difference with the previous snippet is that all inputs have been wrapped in the :py:class:`~aiida.orm.nodes.data.int.Int` class.
 The result that is returned by the function, is now also an :py:class:`~aiida.orm.nodes.data.int.Int` node that can be stored in the provenance graph, and contains the result of the computation.
 
+.. _topics:calculations:concepts:calcfunctions:automatic-serialization:
+
 .. versionadded:: 2.1
 
     If a function argument is a Python base type (i.e. a value of type ``bool``, ``dict``, ``Enum``, ``float``, ``int``, ``list`` or ``str``), it can be passed straight away to the function, without first having to wrap it in the corresponding AiiDA data type.

docs/source/topics/processes/functions.rst

@@ -116,6 +116,57 @@ 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.
 
+Type validation
+===============
+
+.. versionadded:: 2.3
+
+Type hints (introduced with `PEP 484 <https://peps.python.org/pep-0484/>`_ in Python 3.5) can be used to add automatic type validation of process function arguments.
+For example, the following will raise a ``ValueError`` exception:
+
+.. include:: include/snippets/functions/typing_call_raise.py
+    :code: python
+
+When the process function is declared, the process specification (``ProcessSpec``) is built dynamically.
+For each function argument, if a correct type hint is provided, it is set as the ``valid_type`` attribute of the corresponding input port.
+In the example above, the ``x`` and ``y`` inputs have ``Int`` as type hint, which is why the call that passes a ``Float`` raises a ``ValueError``.
+
+.. note::
+
+    Type hints for return values are currently not parsed.
+
+If an argument accepts multiple types, the ``typing.Union`` class can be used as normal:
+
+.. include:: include/snippets/functions/typing_union.py
+    :code: python
+
+The call with an ``Int`` and a ``Float`` will now finish correctly.
+Similarly, optional arguments, with ``None`` as a default, can be declared using ``typing.Optional``:
+
+.. include:: include/snippets/functions/typing_none.py
+    :code: python
+
+The `postponed evaluation of annotations introduced by PEP 563 <https://peps.python.org/pep-0563/>`_ is also supported.
+This means it is possible to use Python base types for the type hint instead of AiiDA's ``Data`` node equivalent:
+
+.. include:: include/snippets/functions/typing_pep_563.py
+    :code: python
+
+The type hints are automatically serialized just as the actual inputs are when the function is called, :ref:`as introduced in v2.1<topics:calculations:concepts:calcfunctions:automatic-serialization>`.
+
+The alternative syntax for union types ``X | Y`` `as introduced by PEP 604 <https://peps.python.org/pep-0604/>`_ is also supported:
+
+.. include:: include/snippets/functions/typing_pep_604.py
+    :code: python
+
+.. warning::
+
+    The usage of notation as defined by PEP 563 and PEP 604 are not supported for Python versions older than 3.10, even if the ``from __future__ import annotations`` statement is added.
+    The reason is that the type inference uses the `inspect.get_annotations <https://docs.python.org/3/library/inspect.html#inspect.get_annotations>`_ method, which was introduced in Python 3.10.
+    For older Python versions, the `get-annotations <https://pypi.org/project/get-annotations>`_ backport is used, but that does not work with PEP 563 and PEP 604, so the constructs from the ``typing`` module have to be used instead.
+
+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.
 
 Return values
 =============