diff --git a/src/aiida/common/__init__.py b/src/aiida/common/__init__.py index 5595ae5849..788894d502 100644 --- a/src/aiida/common/__init__.py +++ b/src/aiida/common/__init__.py @@ -47,6 +47,7 @@ 'FixedFieldsAttributeDict', 'GraphTraversalRule', 'GraphTraversalRules', + 'GraphTraversalRulesType', 'HashingError', 'IncompatibleStorageSchema', 'InputValidationError', diff --git a/src/aiida/common/links.py b/src/aiida/common/links.py index 3043188da3..d3ab228235 100644 --- a/src/aiida/common/links.py +++ b/src/aiida/common/links.py @@ -11,9 +11,11 @@ from collections import namedtuple from enum import Enum +from typing_extensions import TypedDict + from .lang import isidentifier, type_check -__all__ = ('GraphTraversalRule', 'GraphTraversalRules', 'LinkType', 'validate_link_label') +__all__ = ('GraphTraversalRule', 'GraphTraversalRules', 'GraphTraversalRulesType', 'LinkType', 'validate_link_label') class LinkType(Enum): @@ -42,6 +44,23 @@ class LinkType(Enum): """ +class GraphTraversalRulesType(TypedDict, total=False): + """Boolean overrides for graph traversal rules.""" + + input_calc_forward: bool + input_calc_backward: bool + create_forward: bool + create_backward: bool + return_forward: bool + return_backward: bool + input_work_forward: bool + input_work_backward: bool + call_calc_forward: bool + call_calc_backward: bool + call_work_forward: bool + call_work_backward: bool + + class GraphTraversalRules(Enum): """Graph traversal rules when deleting or exporting nodes.""" diff --git a/src/aiida/tools/archive/create.py b/src/aiida/tools/archive/create.py index 51aa57356b..4da79413c3 100644 --- a/src/aiida/tools/archive/create.py +++ b/src/aiida/tools/archive/create.py @@ -14,16 +14,17 @@ import shutil import tempfile -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, cast from tabulate import tabulate +from typing_extensions import Unpack from aiida import orm from aiida.common.lang import type_check -from aiida.common.links import GraphTraversalRules +from aiida.common.links import GraphTraversalRules, GraphTraversalRulesType from aiida.common.log import AIIDA_LOGGER from aiida.common.progress_reporter import get_progress_reporter from aiida.common.utils import DEFAULT_BATCH_SIZE, DEFAULT_FILTER_SIZE, batch_iter @@ -61,7 +62,7 @@ def create_archive( compression: int = 6, test_run: bool = False, backend: StorageBackend | None = None, - **traversal_rules: bool, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> Path: """Export AiiDA data to an archive file. @@ -249,7 +250,7 @@ def querybuilder(): group_nodes, link_data = _collect_required_entities( querybuilder, entity_ids, - traversal_rules, + cast(Mapping[str, bool], traversal_rules), include_authinfos, include_comments, include_logs, @@ -499,7 +500,7 @@ def progress_str(name): def _collect_required_entities( querybuilder: QbType, entity_ids: dict[EntityTypes, set[int]], - traversal_rules: dict[str, bool], + traversal_rules: Mapping[str, bool], include_authinfos: bool, include_comments: bool, include_logs: bool, diff --git a/src/aiida/tools/graph/deletions.py b/src/aiida/tools/graph/deletions.py index 33012ac9f4..d17d650e7d 100644 --- a/src/aiida/tools/graph/deletions.py +++ b/src/aiida/tools/graph/deletions.py @@ -13,6 +13,9 @@ import logging from collections.abc import Callable, Iterable +from typing_extensions import Unpack + +from aiida.common.links import GraphTraversalRulesType from aiida.common.log import AIIDA_LOGGER from aiida.manage import get_manager from aiida.orm import Group, Node, QueryBuilder @@ -28,7 +31,7 @@ def delete_nodes( pks: Iterable[int], dry_run: bool | Callable[[set[int]], bool] = True, backend: StorageBackend | None = None, - **traversal_rules: bool, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> tuple[set[int], bool]: """Delete nodes given a list of "starting" PKs. @@ -114,7 +117,7 @@ def delete_group_nodes( pks: Iterable[int], dry_run: bool | Callable[[set[int]], bool] = True, backend: StorageBackend | None = None, - **traversal_rules: bool, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> tuple[set[int], bool]: """Delete nodes contained in a list of groups (not the groups themselves!). diff --git a/src/aiida/tools/graph/graph_traversers.py b/src/aiida/tools/graph/graph_traversers.py index c51d5c1881..18dd1ee32f 100644 --- a/src/aiida/tools/graph/graph_traversers.py +++ b/src/aiida/tools/graph/graph_traversers.py @@ -13,11 +13,11 @@ from collections.abc import Callable, Iterable from typing import TYPE_CHECKING, Any, cast -from typing_extensions import TypedDict +from typing_extensions import TypedDict, Unpack from aiida import orm from aiida.common import exceptions -from aiida.common.links import GraphTraversalRules, LinkType +from aiida.common.links import GraphTraversalRules, GraphTraversalRulesType, LinkType from aiida.common.progress_reporter import get_progress_reporter from aiida.tools.graph.age_entities import Basket from aiida.tools.graph.age_rules import RuleSaveWalkers, RuleSequence, RuleSetWalkers, UpdateRule @@ -39,7 +39,7 @@ def get_nodes_delete( get_links: bool = False, missing_callback: Callable[[Iterable[int]], None] | None = None, backend: StorageBackend | None = None, - **traversal_rules: bool, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> TraverseGraphOutput: """This function will return the set of all nodes that can be connected to a list of initial nodes through any sequence of specified authorized @@ -79,7 +79,7 @@ def get_nodes_export( starting_pks: Iterable[int], get_links: bool = False, backend: StorageBackend | None = None, - **traversal_rules: bool, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> TraverseGraphOutput: """This function will return the set of all nodes that can be connected to a list of initial nodes through any sequence of specified authorized @@ -116,7 +116,8 @@ def get_nodes_export( def validate_traversal_rules( - ruleset: GraphTraversalRules = GraphTraversalRules.DEFAULT, **traversal_rules: bool + ruleset: GraphTraversalRules = GraphTraversalRules.DEFAULT, + **traversal_rules: Unpack[GraphTraversalRulesType], ) -> dict[str, Any]: """Validates the keywords with a ruleset template and returns a parsed dictionary ready to be used. @@ -135,6 +136,8 @@ def validate_traversal_rules( :param call_work_forward: will traverse CALL_WORK links in the forward direction. :param call_work_backward: will traverse CALL_WORK links in the backward direction. """ + traversal_rules_dict = dict(traversal_rules) + if not isinstance(ruleset, GraphTraversalRules): raise TypeError( f'ruleset input must be of type aiida.common.links.GraphTraversalRules\ninstead, it is: {type(ruleset)}' @@ -147,11 +150,11 @@ def validate_traversal_rules( for name, rule in ruleset.value.items(): follow = rule.default - if name in traversal_rules: + if name in traversal_rules_dict: if not rule.toggleable: raise ValueError(f'input rule {name} is not toggleable for ruleset {ruleset}') - follow = traversal_rules.pop(name) + follow = traversal_rules_dict.pop(name) if not isinstance(follow, bool): raise ValueError(f'the value of rule {name} must be boolean, but it is: {follow}') @@ -166,8 +169,8 @@ def validate_traversal_rules( rules_applied[name] = follow - if traversal_rules: - error_message = f'unrecognized keywords: {", ".join(traversal_rules.keys())}' + if traversal_rules_dict: + error_message = f'unrecognized keywords: {", ".join(traversal_rules_dict.keys())}' raise exceptions.ValidationError(error_message) valid_output = {