Fix typehint of traversal_rules

src/aiida/common/__init__.py

@@ -47,6 +47,7 @@
     'FixedFieldsAttributeDict',
     'GraphTraversalRule',
     'GraphTraversalRules',
+    'GraphTraversalRulesType',
     'HashingError',
     'IncompatibleStorageSchema',
     'InputValidationError',

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."""
 

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,

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!).
 

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 = {