Coverage for core / src / sensorkit / data / filesys.py: 91%
198 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2import asyncio
3import contextlib
4import os
5import pathlib
6from collections.abc import AsyncGenerator
7from fnmatch import fnmatch
8from typing import Literal
10import aiofile
11from loguru import logger
12from pydantic import BaseModel
14from sensorkit.common.aio import cleanup_future
15from sensorkit.common.filewatch import FileEventKind, wait_for_file, watch_dir
16from sensorkit.common.keyword import declare_keyword
17from sensorkit.data.graph import Context, DataFlow, DataOp, SourceOp
18from sensorkit.data.streams import StreamReader, StreamWriter
21@declare_keyword
22class FileNameTemplate(BaseModel):
23 """Template for file naming."""
24 template: str
27@declare_keyword
28class FileInfo(BaseModel):
29 """Information about a file on disk."""
30 path: pathlib.Path
31 size: int | None = None
32 create_time: float | None = None
33 access_time: float | None = None
34 change_time: float | None = None
36 @classmethod
37 async def from_path(cls, path: pathlib.Path, **other):
38 """Create a FileInfo from a local path.
40 Populates: path, size, create_time, access_time, change_time.
41 Additional fields can be provided via keyword arguments.
42 """
43 stat = await asyncio.to_thread(path.stat)
45 return cls.model_construct(
46 path=path,
47 size=stat.st_size,
48 create_time=stat.st_ctime,
49 access_time=stat.st_atime,
50 change_time=stat.st_mtime,
51 **other,
52 )
55async def wait_until_stable(path: pathlib.Path, settle: float) -> bool:
56 """Wait until `path` stops growing, and report whether it still exists.
58 A file becomes visible when the producer creates it, not when it finishes writing it,
59 so reading on first sight yields a truncated result. Polling `st_size` until two reads
60 `settle` apart agree infers the producer is done.
62 This is best-effort and only covers producers that extend a file as they write. It
63 cannot help against one that creates the file at its final size and fills it in place,
64 where the size never changes. Producers that can publish atomically -- write to a
65 temporary name, then rename -- should do so instead and leave `settle` at 0.
67 Returns:
68 False if the file vanished while settling, True otherwise.
69 """
70 if settle <= 0:
71 return True
73 last = -1
75 while True:
76 try:
77 size = (await asyncio.to_thread(path.stat)).st_size
78 except OSError:
79 return False
81 if size == last:
82 return True
84 last = size
85 await asyncio.sleep(settle)
88class WatchDirectory(SourceOp):
89 """DataGraph source that triggers a run for each new file in a directory.
91 Watches a directory with a filesystem observer and, for every newly appeared file
92 whose name matches, starts one graph run seeded with a fresh Context describing that
93 file.
95 Attributes:
96 directory: Path of the directory to watch. Created if it does not exist.
97 match: Glob pattern a filename must match to trigger a run.
98 recursive: Whether to watch subdirectories as well.
100 Supplies context:
101 FileInfo: Names the appeared file. Only `path` is set: the file may still be
102 being written, so stat metadata is left for whoever reads it (`ReadFile`
103 populates it once the contents are settled).
104 """
105 op: Literal["watch_directory"] = "watch_directory"
106 directory: str
107 match: str = "*"
108 recursive: bool = False
110 async def graph_source(self) -> AsyncGenerator[None, DataFlow]:
111 logger.debug(f"starting WatchDirectory source at {self.directory}")
112 await asyncio.to_thread(pathlib.Path(self.directory).mkdir, parents=True, exist_ok=True)
114 async with watch_dir(
115 self.directory,
116 recursive=self.recursive,
117 kinds=(FileEventKind.CREATED, FileEventKind.MOVED),
118 ) as events:
119 while True:
120 # Wait for the graph to be ready.
121 edge = yield
123 # Wait for a matching file to appear.
124 async for event in events:
125 if fnmatch(event.path.name, self.match):
126 path = event.path
127 break
129 logger.debug(f"file {path} appeared")
131 try:
132 await edge.send(
133 Context(FileInfo(path=path)),
134 b"",
135 )
136 except Exception:
137 logger.exception("error sending file path to graph")
140class WriteFile(DataOp):
141 """DataGraph node that writes incoming stream data to a file on disk.
143 If an outgoing edge is connected, the data is forwarded to it as well.
145 Attributes:
146 directory: Base output directory, resolved against context (e.g.
147 "{program_name}"). When None, relative paths and templates resolve under the
148 current working directory.
149 max_chunk_size: Maximum number of bytes read from the stream per chunk.
151 Expects context:
152 FileInfo (optional): An explicit or pre-resolved output path. A relative path is
153 taken under `directory`; an absolute path is used as-is, and is an error if
154 `directory` is also set.
155 FileNameTemplate (optional): A naming template resolved into a name under
156 `directory`, used only when no FileInfo is present.
158 One of FileInfo or FileNameTemplate must be present.
160 Supplies context:
161 FileInfo: Created from FileNameTemplate when no FileInfo was present.
163 Mutates context:
164 FileInfo: Path rewritten to the resolved output location.
166 Raises:
167 ValueError: If neither FileInfo nor FileNameTemplate is present, or if
168 the resolved path is absolute while `directory` is also set.
169 """
170 op: Literal["write_file"] = "write_file"
171 directory: str | None = None
172 max_chunk_size: int = 2**16
174 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
175 assert len(outgoing) <= 1
177 # Read incoming data as a stream.
178 context, reader = await incoming[0].receive("stream")
180 # Resolve the configured base directory against context values (e.g. {program_name}).
181 base = (
182 pathlib.Path(context.resolve(self.directory, as_type=str))
183 if self.directory is not None
184 else pathlib.Path.cwd()
185 )
186 info = context.get(FileInfo)
188 if info is None:
189 # No FileInfo, so look for a FileNameTemplate to resolve the name.
190 template = context.get(FileNameTemplate)
192 if template is None:
193 raise ValueError("WriteFile requires a FileInfo or FileNameTemplate in context")
195 info = FileInfo(path=pathlib.Path(context.resolve(template.template, as_type=str)))
196 context.set(info)
198 # Make sure we don't have two absolute paths.
199 if info.path.is_absolute():
200 if self.directory is not None:
201 raise ValueError(
202 f"ambiguous write location: FileInfo.path is absolute ({info.path}) "
203 f"but WriteFile.directory is also set ({self.directory!r})"
204 )
205 else:
206 # Update the path, applying our base directory.
207 info.path = base / info.path
209 # Ensure the output directory exists before opening the writer.
210 await asyncio.to_thread(info.path.parent.mkdir, parents=True, exist_ok=True)
212 # Get the file writer.
213 writers: list[StreamWriter] = [await get_file_writer(info.path)]
215 if outgoing:
216 # If there is a single outgoing edge, write to that too.
217 writers.append(await outgoing[0].send(context))
219 while not reader.at_eof():
220 chunk = await reader.read(self.max_chunk_size)
222 # FIXME: aiofile requires bytes, which forces us into an unnecessary copy
223 if not isinstance(chunk, bytes):
224 chunk = bytes(chunk)
226 for writer in writers:
227 writer.write(chunk)
229 await asyncio.gather(*(writer.drain() for writer in writers))
230 for writer in writers:
231 writer.close()
233 await asyncio.gather(*(writer.wait_closed() for writer in writers))
236class ReadFile(DataOp):
237 """DataGraph node that reads a file from disk into a downstream stream.
239 Opens the file named by FileInfo.path and streams its contents to the single
240 outgoing edge. The incoming data is expected to be empty.
242 Attributes:
243 max_chunk_size: Maximum number of bytes read from the file per chunk.
244 wait_for_file: Whether to wait for the file to appear before reading.
245 wait_for_file_timeout: Seconds to wait for the file when `wait_for_file` is set.
246 settle_seconds: Interval between the size polls that decide a file still being
247 written has finished. Costs one interval of latency per read. Set to 0 when
248 the producer publishes atomically, or when reading files known to be complete.
250 Expects context:
251 FileInfo: Provides path, the location of the file to read.
253 Mutates context:
254 FileInfo: Replaced with stat metadata (size and timestamps) read from disk when
255 FileInfo.size was not already populated.
256 """
257 op: Literal["read_file"] = "read_file"
258 max_chunk_size: int = 2**16
259 wait_for_file: bool = False
260 wait_for_file_timeout: float = 5.0
261 settle_seconds: float = 0.2
263 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
264 assert len(outgoing) == 1
266 # We expect input data to be empty.
267 context, buffer = await incoming[0].receive("buffer")
268 assert len(buffer) == 0
270 # The path to the file must be present in context.
271 info = context[FileInfo]
273 # Get the file reader.
274 reader = await get_file_reader(
275 info.path,
276 self.max_chunk_size,
277 wait_until_exists=self.wait_for_file_timeout if self.wait_for_file else None,
278 settle=self.settle_seconds,
279 )
281 # Populate stat metadata if it has not been determined yet. The file is known to
282 # exist at this point, since get_file_reader has opened (or waited for) it.
283 if info.size is None:
284 context.set(await FileInfo.from_path(info.path))
286 # Send to the graph.
287 writer = await outgoing[0].send(context)
288 read_count = 0
290 while not reader.at_eof():
291 chunk = await reader.read(self.max_chunk_size)
292 read_count += len(chunk)
293 writer.write(chunk)
295 logger.debug(f"ReadFile read {read_count//1024} KB from {info.path}")
296 writer.close()
297 await writer.wait_closed()
300async def get_file_reader(
301 path: pathlib.Path,
302 read_size: int,
303 wait_until_exists: float | None = None,
304 settle: float = 0.0,
305) -> StreamReader:
306 """Return a StreamReader that asynchronously feeds file content in the background.
308 If *wait_until_exists* is set, waits up to that many seconds for the file to appear.
309 If *settle* is set, waits for the contents to stop growing before opening, so a file
310 caught mid-write is not read truncated (see `wait_until_stable`).
311 """
312 if wait_until_exists is not None:
313 async with asyncio.timeout(wait_until_exists):
314 # Wait until the file exists before opening.
315 await wait_for_file(path)
317 await wait_until_stable(path, settle)
319 ctx = contextlib.AsyncExitStack()
320 f = await ctx.enter_async_context(aiofile.async_open(path, "rb"))
321 reader = asyncio.StreamReader()
323 async def feed_data():
324 async with ctx:
325 while True:
326 chunk = await f.read(read_size)
328 if not chunk:
329 break
331 reader.feed_data(chunk)
333 # Send the EOF signal after the file has closed.
334 reader.feed_eof()
336 # Start feeding data in the background.
337 task = asyncio.create_task(feed_data())
338 task.add_done_callback(cleanup_future)
339 ctx.callback(lambda: task) # Prevent task GC until done
341 return reader
344async def get_file_writer(path: pathlib.Path) -> StreamWriter:
345 """Return a StreamWriter that asynchronously writes to a file, renaming from a .tmp path on close."""
346 return await FileAsyncWriter.open(path)
349class FileAsyncWriter:
350 """Write a file asynchronously using the StreamWriter protocol."""
352 @classmethod
353 async def open(cls, path: pathlib.Path):
354 """Open a new FileAsyncWriter that writes to *path* via a temporary .tmp file."""
355 temp_path = path.with_suffix(".tmp")
356 file = await aiofile.async_open(temp_path, mode="wb")
357 writer = FileAsyncWriter(temp_path, file, path)
358 writer.start()
359 return writer
361 def __init__(
362 self,
363 path: pathlib.Path,
364 file: aiofile.FileIOWrapperBase,
365 target_path: pathlib.Path | None = None,
366 ):
367 self._path = path
368 self._target_path = target_path
369 self._file = file
370 self._queue: asyncio.Queue[bytes] = asyncio.Queue()
371 self._drained = asyncio.Event()
372 self._close_requested = False
374 def start(self):
375 """Start the background write task."""
376 self._task = asyncio.create_task(self._write_task())
378 async def _write_task(self):
379 try:
380 while True:
381 # Get the next data chunk to write and write it.
382 data = await self._queue.get()
383 await self._file.write(data)
384 self._queue.task_done()
386 # Set the drained signal as appropriate.
387 if self._queue.empty():
388 self._drained.set()
389 except Exception:
390 logger.exception("error writing to file")
391 finally:
392 # When we reach here, either:
393 #
394 # 1. The task was cancelled due a call to close() -- this is the happy path.
395 # 2. The task was cancelled externally, e.g. interpreter shutdown.
396 # 3. An error occurred while writing data.
397 #
398 # In all cases we want to try to close the file (which may be partially written in the
399 # error case) and move it to its target path.
400 self._close_requested = True
402 with contextlib.suppress(asyncio.QueueEmpty):
403 while data := self._queue.get_nowait():
404 await self._file.write(data)
405 self._queue.task_done()
407 self._queue.shutdown()
408 self._drained.set()
410 try:
411 async with asyncio.timeout(5.0):
412 logger.debug(f"closing file {self._target_path or self._path}")
413 await self._file.close()
415 if self._target_path:
416 await asyncio.to_thread(
417 os.rename,
418 self._path,
419 self._target_path,
420 )
421 except BaseException as e:
422 final_path = self._target_path or self._path
423 logger.warning(f"error finalizing {final_path}: {str(e)}")
425 def write(self, data: bytes):
426 if self._close_requested:
427 raise RuntimeError("writer is closing or closed")
429 self._queue.put_nowait(data)
430 self._drained.clear()
432 async def drain(self):
433 await self._drained.wait()
435 def close(self):
436 self._close_requested = True
438 # We leverage task cancellation to trigger closing and finalization of the output file.
439 self._task.cancel()
441 def is_closing(self):
442 return self._close_requested and not self._task.done()
444 async def wait_closed(self):
445 with contextlib.suppress(asyncio.CancelledError):
446 await self._task