diff --git a/src/aiida/common/datastructures.py b/src/aiida/common/datastructures.py index 731b1a82b..abd4f2abf 100644 --- a/src/aiida/common/datastructures.py +++ b/src/aiida/common/datastructures.py @@ -141,6 +141,8 @@ class CalcInfo(DefaultFieldsAttributeDict): ) if TYPE_CHECKING: + from aiida.orm.nodes.process.calculation.calcjob import RetrievedList + job_environment: None | dict[str, str] email: None | str email_on_started: bool @@ -154,8 +156,8 @@ class CalcInfo(DefaultFieldsAttributeDict): max_wallclock_seconds: None | int max_memory_kb: None | int rerunnable: bool - retrieve_list: None | list[str | tuple[str, str, str]] - retrieve_temporary_list: None | list[str | tuple[str, str, str]] + retrieve_list: RetrievedList + retrieve_temporary_list: RetrievedList local_copy_list: None | list[tuple[str, str, str]] remote_copy_list: None | list[tuple[str, str, str]] remote_symlink_list: None | list[tuple[str, str, str]] diff --git a/src/aiida/engine/daemon/execmanager.py b/src/aiida/engine/daemon/execmanager.py index ed141bb89..89592727a 100644 --- a/src/aiida/engine/daemon/execmanager.py +++ b/src/aiida/engine/daemon/execmanager.py @@ -20,7 +20,7 @@ from logging import LoggerAdapter from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import TYPE_CHECKING, Any, Optional, Sequence, cast +from typing import TYPE_CHECKING, Any, Optional, cast from typing import Mapping as MappingType from aiida.common import AIIDA_LOGGER, exceptions @@ -35,6 +35,7 @@ from aiida.schedulers.datastructures import JobState if TYPE_CHECKING: + from aiida.orm.nodes.process.calculation.calcjob import RetrievedList from aiida.transports import Transport REMOTE_WORK_DIRECTORY_LOST_FOUND = 'lost+found' @@ -609,7 +610,7 @@ def retrieve_files_from_list( calculation: CalcJobNode, transport: Transport, folder: str, - retrieve_list: Sequence[str | tuple[str, str, int | None]] | None, + retrieve_list: RetrievedList, ) -> None: """Retrieve all the files in the retrieve_list from the remote into the local folder instance through the transport. The entries in the retrieve_list diff --git a/src/aiida/orm/nodes/process/calculation/calcjob.py b/src/aiida/orm/nodes/process/calculation/calcjob.py index e1f7ff735..6e319c2ad 100644 --- a/src/aiida/orm/nodes/process/calculation/calcjob.py +++ b/src/aiida/orm/nodes/process/calculation/calcjob.py @@ -8,8 +8,10 @@ ########################################################################### """Module with `Node` sub class for calculation job processes.""" +from __future__ import annotations + import datetime -from typing import TYPE_CHECKING, Any, AnyStr, Dict, List, Optional, Sequence, Tuple, Type, Union +from typing import TYPE_CHECKING, Any, AnyStr, Dict, List, MutableSequence, Optional, Tuple, Type, Union from aiida.common import exceptions from aiida.common.datastructures import CalcJobState @@ -30,6 +32,8 @@ __all__ = ('CalcJobNode',) +RetrievedList = MutableSequence[str | tuple[str, str, int | None]] | None + class CalcJobNodeCaching(ProcessNodeCaching): """Interface to control caching of a node instance.""" @@ -274,7 +278,7 @@ def get_remote_workdir(self) -> Optional[str]: return self.base.attributes.get(self.REMOTE_WORKDIR_KEY, None) @staticmethod - def _validate_retrieval_directive(directives: Sequence[Union[str, Tuple[str, str, str]]]) -> None: + def _validate_retrieval_directive(directives: RetrievedList) -> None: """Validate a list or tuple of file retrieval directives. :param directives: a list or tuple of file retrieval directives @@ -301,7 +305,7 @@ def _validate_retrieval_directive(directives: Sequence[Union[str, Tuple[str, str if not isinstance(directive[2], (int, type(None))): raise ValueError('invalid directive, third element has to be an integer representing the depth') - def set_retrieve_list(self, retrieve_list: Sequence[Union[str, Tuple[str, str, str]]]) -> None: + def set_retrieve_list(self, retrieve_list: RetrievedList) -> None: """Set the retrieve list. This list of directives will instruct the daemon what files to retrieve after the calculation has completed. @@ -312,14 +316,14 @@ def set_retrieve_list(self, retrieve_list: Sequence[Union[str, Tuple[str, str, s self._validate_retrieval_directive(retrieve_list) self.base.attributes.set(self.RETRIEVE_LIST_KEY, retrieve_list) - def get_retrieve_list(self) -> Sequence[str | tuple[str, str, int | None]] | None: + def get_retrieve_list(self) -> RetrievedList: """Return the list of files/directories to be retrieved on the cluster after the calculation has completed. :return: a list of file directives """ return self.base.attributes.get(self.RETRIEVE_LIST_KEY, None) - def set_retrieve_temporary_list(self, retrieve_temporary_list: Sequence[Union[str, Tuple[str, str, str]]]) -> None: + def set_retrieve_temporary_list(self, retrieve_temporary_list: RetrievedList) -> None: """Set the retrieve temporary list. The retrieve temporary list stores files that are retrieved after completion and made available during parsing @@ -330,7 +334,7 @@ def set_retrieve_temporary_list(self, retrieve_temporary_list: Sequence[Union[st self._validate_retrieval_directive(retrieve_temporary_list) self.base.attributes.set(self.RETRIEVE_TEMPORARY_LIST_KEY, retrieve_temporary_list) - def get_retrieve_temporary_list(self) -> Sequence[str | tuple[str, str, int | None]] | None: + def get_retrieve_temporary_list(self) -> RetrievedList: """Return list of files to be retrieved from the cluster which will be available during parsing. :return: a list of file directives