from functools import partial import io from .abc import AsyncResource from ._util import async_wraps import trio # This list is also in the docs, make sure to keep them in sync _FILE_SYNC_ATTRS = { "closed", "encoding", "errors", "fileno", "isatty", "newlines", "readable", "seekable", "writable", # not defined in *IOBase: "buffer", "raw", "line_buffering", "closefd", "name", "mode", "getvalue", "getbuffer", } # This list is also in the docs, make sure to keep them in sync _FILE_ASYNC_METHODS = { "flush", "read", "read1", "readall", "readinto", "readline", "readlines", "seek", "tell", "truncate", "write", "writelines", # not defined in *IOBase: "readinto1", "peek", } class AsyncIOWrapper(AsyncResource): """A generic :class:`~io.IOBase` wrapper that implements the :term:`asynchronous file object` interface. Wrapped methods that could block are executed in :meth:`trio.to_thread.run_sync`. All properties and methods defined in in :mod:`~io` are exposed by this wrapper, if they exist in the wrapped file object. """ def __init__(self, file): self._wrapped = file @property def wrapped(self): """object: A reference to the wrapped file object""" return self._wrapped def __getattr__(self, name): if name in _FILE_SYNC_ATTRS: return getattr(self._wrapped, name) if name in _FILE_ASYNC_METHODS: meth = getattr(self._wrapped, name) @async_wraps(self.__class__, self._wrapped.__class__, name) async def wrapper(*args, **kwargs): func = partial(meth, *args, **kwargs) return await trio.to_thread.run_sync(func) # cache the generated method setattr(self, name, wrapper) return wrapper raise AttributeError(name) def __dir__(self): attrs = set(super().__dir__()) attrs.update(a for a in _FILE_SYNC_ATTRS if hasattr(self.wrapped, a)) attrs.update(a for a in _FILE_ASYNC_METHODS if hasattr(self.wrapped, a)) return attrs def __aiter__(self): return self async def __anext__(self): line = await self.readline() if line: return line else: raise StopAsyncIteration async def detach(self): """Like :meth:`io.BufferedIOBase.detach`, but async. This also re-wraps the result in a new :term:`asynchronous file object` wrapper. """ raw = await trio.to_thread.run_sync(self._wrapped.detach) return wrap_file(raw) async def aclose(self): """Like :meth:`io.IOBase.close`, but async. This is also shielded from cancellation; if a cancellation scope is cancelled, the wrapped file object will still be safely closed. """ # ensure the underling file is closed during cancellation with trio.CancelScope(shield=True): await trio.to_thread.run_sync(self._wrapped.close) await trio.lowlevel.checkpoint_if_cancelled() async def open_file( file, mode="r", buffering=-1, encoding=None, errors=None, newline=None, closefd=True, opener=None, ): """Asynchronous version of :func:`io.open`. Returns: An :term:`asynchronous file object` Example:: async with await trio.open_file(filename) as f: async for line in f: pass assert f.closed See also: :func:`trio.Path.open` """ _file = wrap_file( await trio.to_thread.run_sync( io.open, file, mode, buffering, encoding, errors, newline, closefd, opener ) ) return _file def wrap_file(file): """This wraps any file object in a wrapper that provides an asynchronous file object interface. Args: file: a :term:`file object` Returns: An :term:`asynchronous file object` that wraps ``file`` Example:: async_file = trio.wrap_file(StringIO('asdf')) assert await async_file.read() == 'asdf' """ def has(attr): return hasattr(file, attr) and callable(getattr(file, attr)) if not (has("close") and (has("read") or has("write"))): raise TypeError( "{} does not implement required duck-file methods: " "close and (read or write)".format(file) ) return AsyncIOWrapper(file)