Skip to content

Commit

Permalink
Add a @pub kwarg to allow specifying a "startup response message"
Browse files Browse the repository at this point in the history
  • Loading branch information
goodboy committed Aug 3, 2021
1 parent f6f9b29 commit bd3a4fd
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion tractor/msg.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ def pub(
wrapped: typing.Callable = None,
*,
tasks: Set[str] = set(),
send_on_connect: Any = None,
):
"""Publisher async generator decorator.
Expand Down Expand Up @@ -206,7 +207,7 @@ async def pub_service(get_topics):

# handle the decorator not called with () case
if wrapped is None:
return partial(pub, tasks=tasks)
return partial(pub, tasks=tasks, send_on_connect=send_on_connect)

task2lock: Dict[str, trio.StrictFIFOLock] = {}

Expand Down Expand Up @@ -249,6 +250,11 @@ async def _execute(

try:
modify_subs(topics2ctxs, topics, ctx)

# if specified send the startup message back to consumer
if send_on_connect is not None:
await ctx.send_yield(send_on_connect)

# block and let existing feed task deliver
# stream data until it is cancelled in which case
# the next waiting task will take over and spawn it again
Expand Down

0 comments on commit bd3a4fd

Please sign in to comment.