Make SharedData generic in type stubs (#106) - #142
Open
aadithyakirantulabandu wants to merge 1 commit into
Open
Conversation
Adds a _SharedData TypeVar to pocketflow/__init__.pyi (bound to Mapping[str, Any] so TypedDict subtypes satisfy it) and threads it through BaseNode/Node/Flow and all subclasses, alongside the existing _PrepResult/_ExecResult/_PostResult TypeVars. Fixes The-Pocket#106: previously `shared` was always typed as the fixed alias `SharedData = Dict[str, Any]`, so mypy/pyright couldn't catch typos or wrong keys even when a caller wanted stricter checking. Now: class MyShared(TypedDict): count: int class Inc(Node[MyShared, None, None, None]): def prep(self, shared: MyShared): shared["cuont"] += 1 # now flagged by the type checker Non-breaking: unparameterized usage (`class Foo(Node): ...`) is unaffected. Only pocketflow/__init__.pyi changes; the runtime pocketflow/__init__.py is untouched, so behavior is identical. Verified with mypy against the stub and by running the full existing test suite locally (33 tests, all passing).
There was a problem hiding this comment.
Pull request overview
This PR updates PocketFlow’s public type stubs to make the shared data structure generic, enabling callers to parameterize nodes/flows with a stricter mapping type (e.g., a TypedDict) for better static checking while keeping runtime behavior unchanged.
Changes:
- Introduces a
_SharedDataTypeVar(bound toMapping[str, Any]) to represent the shared-data shape. - Threads
_SharedDatathroughBaseNode/Node/Flowand async/batch variants soprep/post/runaccept the parameterized shared type. - Keeps
SharedData = Dict[str, Any]as the default alias/documented default shape.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| class BatchFlow(Flow[Optional[List[Params]], Any, _PostResult]): | ||
| def _run(self, shared: SharedData) -> _PostResult: ... | ||
| class BatchFlow(Flow[_SharedData, Optional[List[Params]], Any, _PostResult]): |
| async def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ... | ||
|
|
||
| class AsyncFlow(Flow[_PrepResult, Any, _PostResult], AsyncNode[_PrepResult, Any, _PostResult]): | ||
| class AsyncFlow(Flow[_SharedData, _PrepResult, Any, _PostResult], AsyncNode[_SharedData, _PrepResult, Any, _PostResult]): |
|
|
||
| class AsyncBatchFlow(AsyncFlow[Optional[List[Params]], Any, _PostResult], BatchFlow[Optional[List[Params]], Any, _PostResult]): | ||
| async def _run_async(self, shared: SharedData) -> _PostResult: ... | ||
| class AsyncBatchFlow(AsyncFlow[_SharedData, Optional[List[Params]], Any, _PostResult], BatchFlow[_SharedData, Optional[List[Params]], Any, _PostResult]): |
|
|
||
| class AsyncParallelBatchFlow(AsyncFlow[Optional[List[Params]], Any, _PostResult], BatchFlow[Optional[List[Params]], Any, _PostResult]): | ||
| async def _run_async(self, shared: SharedData) -> _PostResult: ... No newline at end of file | ||
| class AsyncParallelBatchFlow(AsyncFlow[_SharedData, Optional[List[Params]], Any, _PostResult], BatchFlow[_SharedData, Optional[List[Params]], Any, _PostResult]): |
| class BaseNode(Generic[_SharedData, _PrepResult, _ExecResult, _PostResult]): | ||
| params: Params | ||
| successors: Dict[str, BaseNode[Any, Any, Any]] | ||
| successors: Dict[str, BaseNode[Any, Any, Any, Any]] |
| def set_params(self, params: Params) -> None: ... | ||
| def next(self, node: BaseNode[Any, Any, Any], action: str = "default") -> BaseNode[Any, Any, Any]: ... | ||
| def prep(self, shared: SharedData) -> _PrepResult: ... | ||
| def next(self, node: BaseNode[Any, Any, Any, Any], action: str = "default") -> BaseNode[Any, Any, Any, Any]: ... |
Comment on lines
+29
to
30
| def __rshift__(self, other: BaseNode[Any, Any, Any, Any]) -> BaseNode[Any, Any, Any, Any]: ... | ||
| def __sub__(self, action: str) -> _ConditionalTransition: ... |
Comment on lines
32
to
+37
| class _ConditionalTransition: | ||
| src: BaseNode[Any, Any, Any] | ||
| src: BaseNode[Any, Any, Any, Any] | ||
| action: str | ||
|
|
||
| def __init__(self, src: BaseNode[Any, Any, Any], action: str) -> None: ... | ||
| def __rshift__(self, tgt: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ... | ||
| def __init__(self, src: BaseNode[Any, Any, Any, Any], action: str) -> None: ... | ||
| def __rshift__(self, tgt: BaseNode[Any, Any, Any, Any]) -> BaseNode[Any, Any, Any, Any]: ... |
Comment on lines
+52
to
+58
| start_node: Optional[BaseNode[Any, Any, Any, Any]] | ||
|
|
||
| def __init__(self, start: Optional[BaseNode[Any, Any, Any]] = None) -> None: ... | ||
| def start(self, start: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ... | ||
| def __init__(self, start: Optional[BaseNode[Any, Any, Any, Any]] = None) -> None: ... | ||
| def start(self, start: BaseNode[Any, Any, Any, Any]) -> BaseNode[Any, Any, Any, Any]: ... | ||
| def get_next_node( | ||
| self, curr: BaseNode[Any, Any, Any], action: Optional[str] | ||
| ) -> Optional[BaseNode[Any, Any, Any]]: ... | ||
| self, curr: BaseNode[Any, Any, Any, Any], action: Optional[str] | ||
| ) -> Optional[BaseNode[Any, Any, Any, Any]]: ... |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a _SharedData TypeVar to pocketflow/init.pyi (bound to Mapping[str, Any] so TypedDict subtypes satisfy it) and threads it through BaseNode/Node/Flow and all subclasses, alongside the existing _PrepResult/_ExecResult/_PostResult TypeVars.
Fixes #106: previously
sharedwas always typed as the fixed aliasSharedData = Dict[str, Any], so mypy/pyright couldn't catch typos or wrong keys even when a caller wanted stricter checking. Now:Non-breaking: unparameterized usage (
class Foo(Node): ...) is unaffected. Only pocketflow/init.pyi changes; the runtime pocketflow/init.py is untouched, so behavior is identical. Verified with mypy against the stub and by running the full existing test suite locally (33 tests, all passing).