Skip to content

Make SharedData generic in type stubs (#106) #142

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aadithyakirantulabandu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make SharedData generic in type stubs (#106) #142

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 44 additions & 42 deletions pocketflow/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,40 +1,42 @@
import asyncio
from typing import Any, Dict, List, Optional, Union, TypeVar, Generic
from typing import Any, Dict, List, Mapping, Optional, Union, TypeVar, Generic

# Type variables for better type relationships
_PrepResult = TypeVar('_PrepResult')
_ExecResult = TypeVar('_ExecResult')
_PostResult = TypeVar('_PostResult')
# Bound to Mapping (not Dict) so structural types like TypedDict satisfy it too.
_SharedData = TypeVar('_SharedData', bound=Mapping[str, Any])

# More specific parameter types
ParamValue = Union[str, int, float, bool, None, List[Any], Dict[str, Any]]
SharedData = Dict[str, Any]
SharedData = Dict[str, Any] # default shape; parametrize a node/flow with your own TypedDict for stricter checking
Params = Dict[str, ParamValue]

class BaseNode(Generic[_PrepResult, _ExecResult, _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 __init__(self) -> None: ...
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]: ...
def prep(self, shared: _SharedData) -> _PrepResult: ...
def exec(self, prep_res: _PrepResult) -> _ExecResult: ...
def post(self, shared: SharedData, prep_res: _PrepResult, exec_res: _ExecResult) -> _PostResult: ...
def post(self, shared: _SharedData, prep_res: _PrepResult, exec_res: _ExecResult) -> _PostResult: ...
def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...
def _run(self, shared: SharedData) -> _PostResult: ...
def run(self, shared: SharedData) -> _PostResult: ...
def __rshift__(self, other: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ...
def _run(self, shared: _SharedData) -> _PostResult: ...
def run(self, shared: _SharedData) -> _PostResult: ...
def __rshift__(self, other: BaseNode[Any, Any, Any, Any]) -> BaseNode[Any, Any, Any, Any]: ...
def __sub__(self, action: str) -> _ConditionalTransition: ...
Comment on lines +29 to 30

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 32 to +37

class Node(BaseNode[_PrepResult, _ExecResult, _PostResult]):
class Node(BaseNode[_SharedData, _PrepResult, _ExecResult, _PostResult]):
max_retries: int
wait: Union[int, float]
cur_retry: int
Expand All @@ -43,55 +45,55 @@ class Node(BaseNode[_PrepResult, _ExecResult, _PostResult]):
def exec_fallback(self, prep_res: _PrepResult, exc: Exception) -> _ExecResult: ...
def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...

class BatchNode(Node[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
class BatchNode(Node[_SharedData, Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ...

class Flow(BaseNode[_PrepResult, Any, _PostResult]):
start_node: Optional[BaseNode[Any, Any, Any]]
class Flow(BaseNode[_SharedData, _PrepResult, Any, _PostResult]):
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]]: ...
Comment on lines +52 to +58
def _orch(
self, shared: SharedData, params: Optional[Params] = None
self, shared: _SharedData, params: Optional[Params] = None
) -> Any: ...
def _run(self, shared: SharedData) -> _PostResult: ...
def post(self, shared: SharedData, prep_res: _PrepResult, exec_res: Any) -> _PostResult: ...
def _run(self, shared: _SharedData) -> _PostResult: ...
def post(self, shared: _SharedData, prep_res: _PrepResult, exec_res: Any) -> _PostResult: ...

class BatchFlow(Flow[Optional[List[Params]], Any, _PostResult]):
def _run(self, shared: SharedData) -> _PostResult: ...
class BatchFlow(Flow[_SharedData, Optional[List[Params]], Any, _PostResult]):
def _run(self, shared: _SharedData) -> _PostResult: ...

class AsyncNode(Node[_PrepResult, _ExecResult, _PostResult]):
async def prep_async(self, shared: SharedData) -> _PrepResult: ...
class AsyncNode(Node[_SharedData, _PrepResult, _ExecResult, _PostResult]):
async def prep_async(self, shared: _SharedData) -> _PrepResult: ...
async def exec_async(self, prep_res: _PrepResult) -> _ExecResult: ...
async def exec_fallback_async(self, prep_res: _PrepResult, exc: Exception) -> _ExecResult: ...
async def post_async(
self, shared: SharedData, prep_res: _PrepResult, exec_res: _ExecResult
self, shared: _SharedData, prep_res: _PrepResult, exec_res: _ExecResult
) -> _PostResult: ...
async def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...
async def run_async(self, shared: SharedData) -> _PostResult: ...
async def _run_async(self, shared: SharedData) -> _PostResult: ...
def _run(self, shared: SharedData) -> _PostResult: ...
async def run_async(self, shared: _SharedData) -> _PostResult: ...
async def _run_async(self, shared: _SharedData) -> _PostResult: ...
def _run(self, shared: _SharedData) -> _PostResult: ...

class AsyncBatchNode(AsyncNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
class AsyncBatchNode(AsyncNode[_SharedData, Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[_SharedData, Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
async def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ...

class AsyncParallelBatchNode(AsyncNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
class AsyncParallelBatchNode(AsyncNode[_SharedData, Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[_SharedData, Optional[List[_PrepResult]], List[_ExecResult], _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]):
async def _orch_async(
self, shared: SharedData, params: Optional[Params] = None
self, shared: _SharedData, params: Optional[Params] = None
) -> Any: ...
async def _run_async(self, shared: SharedData) -> _PostResult: ...
async def _run_async(self, shared: _SharedData) -> _PostResult: ...
async def post_async(
self, shared: SharedData, prep_res: _PrepResult, exec_res: Any
self, shared: _SharedData, prep_res: _PrepResult, exec_res: 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]):
async def _run_async(self, shared: _SharedData) -> _PostResult: ...

class AsyncParallelBatchFlow(AsyncFlow[Optional[List[Params]], Any, _PostResult], BatchFlow[Optional[List[Params]], Any, _PostResult]):
async def _run_async(self, shared: SharedData) -> _PostResult: ...
class AsyncParallelBatchFlow(AsyncFlow[_SharedData, Optional[List[Params]], Any, _PostResult], BatchFlow[_SharedData, Optional[List[Params]], Any, _PostResult]):
async def _run_async(self, shared: _SharedData) -> _PostResult: ...