Skip to content

Commit ff4be5e

Browse files
committed
Provide a nicer, type-safe way of adding calls to the batch
1 parent 81a2df6 commit ff4be5e

File tree

2 files changed

+8
-4
lines changed

2 files changed

+8
-4
lines changed

src/dispatch/function.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -359,8 +359,12 @@ def __init__(self, client: Client):
359359
self.client = client
360360
self.calls: list[Call] = []
361361

362-
def add(self, call: Call) -> Batch:
363-
"""Add a call to the batch."""
362+
def add(self, func: Function[P, T], *args: P.args, **kwargs: P.kwargs) -> Batch:
363+
"""Add a call to the specified function to the batch."""
364+
return self.add_call(func.build_call(*args, correlation_id=None, **kwargs))
365+
366+
def add_call(self, call: Call) -> Batch:
367+
"""Add a Call to the batch."""
364368
self.calls.append(call)
365369
return self
366370

tests/test_client.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,9 @@ def test_call_pickle(self):
7575
def test_batch(self):
7676
batch = self.dispatch_client.batch()
7777

78-
batch.add(Call(function="my-function", input=42)).add(
78+
batch.add_call(Call(function="my-function", input=42)).add_call(
7979
Call(function="my-function", input=23)
80-
).add(Call(function="my-function2", input=11))
80+
).add_call(Call(function="my-function2", input=11))
8181

8282
dispatch_ids = batch.dispatch()
8383
self.assertEqual(len(dispatch_ids), 3)

0 commit comments

Comments
 (0)