1 27
from functools import partial
2 27
import io
3

4 27
from .abc import AsyncResource
5 27
from ._util import async_wraps
6

7 27
import trio
8

9
# This list is also in the docs, make sure to keep them in sync
10 27
_FILE_SYNC_ATTRS = {
11
    "closed",
12
    "encoding",
13
    "errors",
14
    "fileno",
15
    "isatty",
16
    "newlines",
17
    "readable",
18
    "seekable",
19
    "writable",
20
    # not defined in *IOBase:
21
    "buffer",
22
    "raw",
23
    "line_buffering",
24
    "closefd",
25
    "name",
26
    "mode",
27
    "getvalue",
28
    "getbuffer",
29
}
30

31
# This list is also in the docs, make sure to keep them in sync
32 27
_FILE_ASYNC_METHODS = {
33
    "flush",
34
    "read",
35
    "read1",
36
    "readall",
37
    "readinto",
38
    "readline",
39
    "readlines",
40
    "seek",
41
    "tell",
42
    "truncate",
43
    "write",
44
    "writelines",
45
    # not defined in *IOBase:
46
    "readinto1",
47
    "peek",
48
}
49

50

51 27
class AsyncIOWrapper(AsyncResource):
52
    """A generic :class:`~io.IOBase` wrapper that implements the :term:`asynchronous
53
    file object` interface. Wrapped methods that could block are executed in
54
    :meth:`trio.to_thread.run_sync`.
55

56
    All properties and methods defined in in :mod:`~io` are exposed by this
57
    wrapper, if they exist in the wrapped file object.
58

59
    """
60

61 27
    def __init__(self, file):
62 27
        self._wrapped = file
63

64 27
    @property
65 13
    def wrapped(self):
66
        """object: A reference to the wrapped file object"""
67

68 27
        return self._wrapped
69

70 27
    def __getattr__(self, name):
71 27
        if name in _FILE_SYNC_ATTRS:
72 27
            return getattr(self._wrapped, name)
73 27
        if name in _FILE_ASYNC_METHODS:
74 27
            meth = getattr(self._wrapped, name)
75

76 27
            @async_wraps(self.__class__, self._wrapped.__class__, name)
77 13
            async def wrapper(*args, **kwargs):
78 27
                func = partial(meth, *args, **kwargs)
79 27
                return await trio.to_thread.run_sync(func)
80

81
            # cache the generated method
82 27
            setattr(self, name, wrapper)
83 27
            return wrapper
84

85 27
        raise AttributeError(name)
86

87 27
    def __dir__(self):
88 27
        attrs = set(super().__dir__())
89 27
        attrs.update(a for a in _FILE_SYNC_ATTRS if hasattr(self.wrapped, a))
90 27
        attrs.update(a for a in _FILE_ASYNC_METHODS if hasattr(self.wrapped, a))
91 27
        return attrs
92

93 27
    def __aiter__(self):
94 27
        return self
95

96 27
    async def __anext__(self):
97 27
        line = await self.readline()
98 27
        if line:
99 27
            return line
100
        else:
101 27
            raise StopAsyncIteration
102

103 27
    async def detach(self):
104
        """Like :meth:`io.BufferedIOBase.detach`, but async.
105

106
        This also re-wraps the result in a new :term:`asynchronous file object`
107
        wrapper.
108

109
        """
110

111 27
        raw = await trio.to_thread.run_sync(self._wrapped.detach)
112 27
        return wrap_file(raw)
113

114 27
    async def aclose(self):
115
        """Like :meth:`io.IOBase.close`, but async.
116

117
        This is also shielded from cancellation; if a cancellation scope is
118
        cancelled, the wrapped file object will still be safely closed.
119

120
        """
121

122
        # ensure the underling file is closed during cancellation
123 27
        with trio.CancelScope(shield=True):
124 27
            await trio.to_thread.run_sync(self._wrapped.close)
125

126 27
        await trio.lowlevel.checkpoint_if_cancelled()
127

128

129 27
async def open_file(
130
    file,
131
    mode="r",
132
    buffering=-1,
133
    encoding=None,
134
    errors=None,
135
    newline=None,
136
    closefd=True,
137
    opener=None,
138
):
139
    """Asynchronous version of :func:`io.open`.
140

141
    Returns:
142
        An :term:`asynchronous file object`
143

144
    Example::
145

146
        async with await trio.open_file(filename) as f:
147
            async for line in f:
148
                pass
149

150
        assert f.closed
151

152
    See also:
153
      :func:`trio.Path.open`
154

155
    """
156 27
    _file = wrap_file(
157
        await trio.to_thread.run_sync(
158
            io.open, file, mode, buffering, encoding, errors, newline, closefd, opener
159
        )
160
    )
161 27
    return _file
162

163

164 27
def wrap_file(file):
165
    """This wraps any file object in a wrapper that provides an asynchronous
166
    file object interface.
167

168
    Args:
169
        file: a :term:`file object`
170

171
    Returns:
172
        An :term:`asynchronous file object` that wraps ``file``
173

174
    Example::
175

176
        async_file = trio.wrap_file(StringIO('asdf'))
177

178
        assert await async_file.read() == 'asdf'
179

180
    """
181

182 27
    def has(attr):
183 27
        return hasattr(file, attr) and callable(getattr(file, attr))
184

185 27
    if not (has("close") and (has("read") or has("write"))):
186 27
        raise TypeError(
187
            "{} does not implement required duck-file methods: "
188
            "close and (read or write)".format(file)
189
        )
190

191 27
    return AsyncIOWrapper(file)

Read our documentation on viewing source code .

Loading