This commit is contained in:
Eugene Yurtsev
2025-03-19 16:08:25 -04:00
parent 73f9ef0ef8
commit c9cb8165d4
3 changed files with 8 additions and 8 deletions
+2 -2
View File
@@ -2751,10 +2751,10 @@ class Pregel(PregelProtocol):
return chunks
def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> dict[str, list[str]]:
def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, list[str]]:
"""Index from a trigger to nodes that depend on it."""
trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list)
for name, node in nodes.items():
for trigger in node.triggers:
trigger_to_nodes.setdefault(trigger, []).append(name)
return trigger_to_nodes
return cast(Mapping[str, list[str]], trigger_to_nodes)
+3 -3
View File
@@ -347,7 +347,7 @@ def prepare_next_tasks(
store: Literal[None] = None,
checkpointer: Literal[None] = None,
manager: Literal[None] = None,
trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
updated_channels: Optional[set[str]] = None,
) -> dict[str, PregelTask]: ...
@@ -366,7 +366,7 @@ def prepare_next_tasks(
store: Optional[BaseStore],
checkpointer: Optional[BaseCheckpointSaver],
manager: Union[None, ParentRunManager, AsyncParentRunManager],
trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
updated_channels: Optional[set[str]] = None,
) -> dict[str, PregelExecutableTask]: ...
@@ -384,7 +384,7 @@ def prepare_next_tasks(
store: Optional[BaseStore] = None,
checkpointer: Optional[BaseCheckpointSaver] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
updated_channels: Optional[set[str]] = None,
) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
"""Prepare the set of tasks that will make up the next Pregel step.
+3 -3
View File
@@ -209,7 +209,7 @@ class PregelLoop(LoopProtocol):
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
input_model: Optional[Type[BaseModel]] = None,
debug: bool = False,
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
) -> None:
super().__init__(
step=0,
@@ -889,7 +889,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
input_model: Optional[Type[BaseModel]] = None,
debug: bool = False,
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
) -> None:
super().__init__(
input,
@@ -1032,7 +1032,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
input_model: Optional[Type[BaseModel]] = None,
debug: bool = False,
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
) -> None:
super().__init__(
input,