Coverage for core / src / sensorkit / data / local.py: 100%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5from collections.abc import Buffer 

6from typing import Any, AsyncGenerator, Literal 

7 

8from loguru import logger 

9 

10from sensorkit.data.context import Context 

11from sensorkit.data.graph import DataFlow, SinkOp, SourceOp 

12from sensorkit.data.streams import StreamReader, StreamWriter, create_connected_streams 

13 

14 

15class AppSource(SourceOp): 

16 """Source node for a DataGraph.""" 

17 op: Literal["app_source"] = "app_source" 

18 

19 def __init__(self, /, **data: Any): 

20 super().__init__(**data) 

21 self._queue: asyncio.Queue[tuple[Context, Buffer | StreamReader]] = asyncio.Queue() 

22 

23 def produce( 

24 self, 

25 context: Context | None = None, 

26 obj: Buffer | StreamReader | None = None, 

27 ) -> StreamWriter | None: 

28 """Enqueue a graph run. 

29 

30 If *obj* is omitted, a connected StreamWriter is returned for the caller to write into. 

31 """ 

32 result = None 

33 

34 if obj is None: 

35 obj, result = create_connected_streams() 

36 

37 self._queue.put_nowait((context or Context(), obj)) 

38 return result 

39 

40 async def graph_source(self) -> AsyncGenerator[None, DataFlow]: 

41 while True: 

42 edge = yield 

43 context, obj = await self._queue.get() 

44 

45 try: 

46 logger.debug(f"AppSource executing DataGraph with {context=}") 

47 await edge.send(context, obj) 

48 finally: 

49 self._queue.task_done() 

50 

51 

52class AppSink(SinkOp): 

53 """Sink node for a DataGraph.""" 

54 op: Literal["app_sink"] = "app_sink" 

55 

56 def __init__(self, /, **data: Any): 

57 super().__init__(**data) 

58 self._queue: asyncio.Queue[tuple[Context, Buffer]] = asyncio.Queue() 

59 

60 async def consume(self): 

61 """Async generator that yields `(context, buffer)` tuples as graph runs complete.""" 

62 while True: 

63 record = await self._queue.get() 

64 self._queue.task_done() 

65 yield record 

66 

67 async def graph_sink(self, incoming: DataFlow): 

68 record = await incoming.receive("buffer") 

69 await self._queue.put(record) 

70 logger.debug(f"AppSink consumed DataGraph with context={record[0]}")