1# ruff: noqa: PLW0603,PLW2901 # ty: ignore[unresolved-attribute]
2from asyncutils._internal import helpers as H, patch as P
3from asyncutils._internal.submodules import iotools_all as __all__
4import asyncutils as A, errno as E, os as O, sys as S
5from _functools import partial
6from asyncio import Lock, gather
7from contextlib import asynccontextmanager
8from itertools import count, starmap
9from mmap import mmap
10def f(a, b, f=S.audit, _=O.pipe, /):
11 def double_ended_pipe(*, pipe_impl=_, x=partial(open, mode=a), y=partial(open, mode=b), f=f): r, W, R, w = *pipe_impl(), *pipe_impl(); f(f'asyncutils.iotools.double_ended_{'text' if a == 'r' else 'binary'}_pipe', r, w, R, W); return tuple(map(AsyncReadWriteCouple, map(x, (r, R)), map(y, (w, W))))
12 return double_ended_pipe
13z, s = lambda s=None, /, **d: {k: v for k, v in d.items() if v is not s}, '*, pipe_impl={}'
14double_ended_text_pipe, double_ended_binary_pipe = t = tuple(map(f, ('r', 'rb'), ('w', 'wb')))
15P.patch_function_signatures(*((_, s) for _ in t))
[docs]
16@H.subscriptable
17class AsyncReadWriteCouple(A.LoopContextMixin):
18 __slots__ = '__aso', '__x', 'reader', 'writer'; executor = None
19 def __init__(self, r, w, /, executor=None, _=H.create_executor, *, find_attr_on_writer_first=False):
20 if not r.readable(): raise TypeError(f'asyncutils.iotools.AsyncReadWriteCouple: reader {r!r} is not readable')
21 if not w.writable(): raise TypeError(f'asyncutils.iotools.AsyncReadWriteCouple: writer {w!r} is not writable')
22 super().__init__(); self.__x, self.reader, self.writer, self.__aso = _(t) if executor is None is (executor := (t := type(self)).executor) else executor, r, w, (w, r) if find_attr_on_writer_first else (r, w)
23 async def _run(self, f, *a): return await self.loop.run_in_executor(self.__x, f, *a)
[docs]
24 def read(self, n=-1, /): return self._run(self.reader.read, n)
[docs]
25 def read1(self, n=-1, /): return self._run(self.reader.read1, n)
[docs]
26 def readall(self): return self._run(r.read if (f := getattr(r := self.reader, 'readall', None)) is None else f)
[docs]
27 async def readinto(self, b, /):
28 if (f := getattr(r := self.reader, 'readinto', None)) is not None: return await self._run(f, b)
29 if (m := memoryview(b)).readonly: raise TypeError('asyncutils.iotools.AsyncReadWriteCouple: cannot read into a read-only buffer')
30 if l := len(d := await self._run(r.read, m.nbytes)): m[:l] = d # cspell:disable-line
31 return l
[docs]
32 def readinto1(self, b, /): return self._run(self.reader.readinto1, b)
[docs]
33 def readline(self, limit=-1, /): return self._run(self.reader.readline, limit)
[docs]
34 def readlines(self, hint=-1, /): return self._run(self.reader.readlines, hint)
[docs]
35 def write(self, s, /): return self._run(self.writer.write, s)
[docs]
36 def writelines(self, lines, /): return self._run(self.writer.writelines, lines)
37 def __rx(self, e, m, t='tip: delegate to reader or writer as appropriate', /): A.raise_exc(OSError, e, f'asyncutils.iotools.AsyncReadWriteCouple at {id(self):#x}: ambiguous {m} call', notes=t)
[docs]
38 def fileno(self): self.__rx(E.EBADF, 'fileno')
[docs]
39 def isatty(self): self.__rx(E.ENOTSUP, 'isatty')
40 readable = writable = lambda _, /: True
[docs]
41 def flush(self): return self._run(self.writer.flush)
[docs]
42 def seekable(self): return False # noqa: PLR6301
[docs]
43 def seek(self, *_): self.__rx(E.ESPIPE, 'seek')
[docs]
44 def tell(self): self.__rx(E.ESPIPE, 'tell')
[docs]
45 def truncate(self, size=None, /): return self._run(self.writer.truncate, size)
[docs]
46 async def aclose(self): await gather(*map(self._run, (self.reader.close, self.writer.close))); self.__x.shutdown()
47 __cleanup__ = aclose
48 @property
49 def closed(self): return self.reader.closed and self.writer.closed
[docs]
50 def __getattr__(self, n, /):
51 f = (a := []).append
52 for _ in self.__aso:
53 try: return getattr(_, n)
54 except AttributeError as e: f(e)
55 raise ExceptionGroup(f'asyncutils.iotools.AsyncReadWriteCouple: did not find attribute {n!r}', a) from None
56class File(A.LoopContextMixin): # noqa: PLR0904
57 __slots__ = '__f', '__m', '__n'
58 if S.platform != 'win32':
59 def madvise(self, option, start=0, length=None, _=H.filter_out): return self.__m.madvise(option, start, *_(length)) # ty: ignore[possibly-missing-attribute]
60 def read(self, offset=0, size=-1): return self._run(self._read, offset, size)
61 def write(self, data, offset=0): return self._run(self._write, data, offset)
62 async def readline(self, offset=0, size=None, include_newline=False): return (await self._run(self._readline, offset, size, include_newline))[0]
63 async def readlines(self, hint=-1): return list(await self._run(self._readlines, hint))
64 async def flush(self, offset=0, size=None, /): return await self._run(self._flush, offset, size)
65 def move(self, dest, src, count): return self._run(self.__m.move, dest, src, count)
66 async def __setup__(self): self.__m = m = mmap(self.__n, 0, access=2).__enter__(); self.mgr.add(m)
67 async def __cleanup__(self): await self.aclose(); self.mgr.discard(self.__m)
68 def seek(self, pos, whence=0): return self._run(self.__m.seek, pos, whence)
69 def __new__(cls, file, /):
70 if (r := (f := cls.open_files.get)((file, 'r+b'))) is None is (r := f((file, 'w+b'))) is (r := f((file, 'x+b'))): (r := super().__new__(cls)).__f, r.__n = file, file.fileno()
71 return r
72 def __iter__(self): return self.__f.__iter__()
73 def __aiter__(self): return A.iter_to_agen(self.__f)
74 def __del__(self): self.make(self.aclose())
75 @property
76 def closed(self): return self.__f.closed
77 def fileno(self): return self.__n
78 def sync(self, _=O.fsync): self._flush(0, None); _(self.__n)
79 async def aclose(self): await gather(*map(self._run, (self.__m.close, self.__f.close)))
80 def close(self): self.__m.close(); self.__f.close()
81 def read_byte(self): return self.__m.read_byte()
82 def write_byte(self, b, /): self.__m.write_byte(b)
83 def resize(self, new_size): self.__m.resize(new_size)
84 def find(self, sub, start=None, end=None, _=z): return self.__m.find(sub, **_(start=start, end=end))
85 def rfind(self, sub, start=None, end=None, _=z): return self.__m.rfind(sub, **_(start=start, end=end))
86 def tell(self): return self.__m.tell()
87 def size(self): return self.__m.size()
88 def isatty(self): return self.__f.isatty()
89 readable = writable = seekable = AsyncReadWriteCouple.readable
90 def _flush(self, offset, size, _=H.filter_out): self.__f.flush(); self.__m.flush(offset, *_(size))
91 def _trunc_from(self, data, offset): c = (m := self.__m).tell(); m.seek(0, 2); m.resize(max(m.tell(), x := offset+len(data))); m.seek(c); return x
92 def _read(self, offset, size): return self.__m[offset:None if size < 0 else offset+size]
93 def _write(self, data, offset): (m := self.__m)[offset:self._trunc_from(data, offset)] = data; m.flush()
94 def _readline(self, offset, size, include_newline): return (b'', 0) if offset >= (l := len(m := self.__m)) else (m[offset:(q := p if (e := m.find(b'\n', offset, p := (l if size is None else min(offset+size, l)))) == -1 else e+include_newline)], q)
95 def _readlines(self, hint, /):
96 if hint < 0: yield from map(bytes, self.__m); return
97 f = self._readline
98 while hint > 0: b, n = f(0, None, False); yield b; hint -= n
99 async def writelines(self, lines, /, *, sep=b'', minimize_writes=None):
100 f, lines = self.write, A.iter_to_agen(lines)
101 if A.getcontext().MEMORY_MAPPED_IO_MANAGER_DEFAULT_MINIMIZE_WRITES if minimize_writes is None else minimize_writes: return await f(sep.join(await A.to_list(lines)))
102 if sep:
103 async for l in lines: await f(l); await f(sep)
104 else:
105 async for l in lines: await f(l)
106 async def read_str(self, offset=0, size=-1, encoding='utf-8', errors='strict'): return (await self.read(offset, size)).decode(encoding, errors)
107 def write_str(self, text, offset=0, encoding='utf-8', errors='strict'): return self.write(text.encode(encoding, errors), offset)
108 def smart_write(self, data, offset=0, encoding='utf-8', errors='strict'): return self.write(data.encode(encoding, errors) if isinstance(data, str) else data, offset)
109 async def copy_range(self, src_offset, dest_offset, size):
110 try: await self.write(await self.read(src_offset, size), dest_offset); return True
111 except: return False # noqa: E722
112 def fill(self, pattern, offset=0, count=1): return self.write(pattern*count, offset)
113 async def compare(self, other, /, size=-1, offset_self=0, offset_other=0): return (await self.read(offset_self, size)) == (await other.read(offset_other, size))
114 async def hamming_dist(self, other, /, size=-1, offset_self=0, offset_other=0, _=tuple(map(int.bit_count, range(0x100)))): return sum(_[i^j] for i, j in zip(await self.read(offset_self, size), await other.read(offset_other, size), strict=size > 0)) # noqa: B008
115 async def hamming_dist_bytes(self, other, /, size=-1, offset_self=0, offset_other=0): return sum(i != j for i, j in zip(await self.read(offset_self, size), await other.read(offset_other, size), strict=size > 0))
116 async def read_until(self, delim, offset=0, maxsize=-1): return (d, offset+len(d)) if (p := (d := await self.read(offset, maxsize)).find(delim)) == -1 else (d[:p+(l := len(delim))], offset+p+l)
117 async def insert(self, data, offset): await self.write(data if offset > await self._run(self.size) else data+await self.read(offset), offset)
118 async def delete(self, offset, size):
119 if size <= 0 or offset >= (s := await self._run(self.__m.size)): return
120 if (t := offset+size) < s: await self.write(await self.read(t), offset)
121 await self._run(self.resize, max(0, s-size))
122 async def replace(self, old, new, offset=0, count=None):
123 r, c, o, n, f, g, h = 0, offset, len(old), len(new), partial(self._run, self.find, old), self.delete, self.insert
124 if count is None: count = float('inf')
125 while r < count:
126 if (p := await f(c)) == -1: break
127 await g(p, o); await h(new, p); r += 1; c = p+n
128 return r
129 async def search_lazy(self, pattern, offset=0):
130 f = partial(self._run, self.find, pattern)
131 for c in count(offset):
132 if (p := await f(c)) == -1: break
133 yield p
134 async def search_lazy_non_overlapping(self, pattern, offset=0):
135 f = partial(self._run, self.find, pattern)
136 while True:
137 if (offset := await f(offset)) == -1: break
138 yield offset
139 def search(self, pattern, offset=0, max_results=None): return A.collect(self.search_lazy(pattern, offset), max_results)
140 def search_non_overlapping(self, pattern, offset=0, max_results=None): return A.collect(self.search_lazy_non_overlapping(pattern, offset), max_results)
141 async def compact(self):
142 for i in range(len(c := await self.read())):
143 if c[~i]: await self._run(self.resize, c-i); return i
144 def __init_subclass__(cls, *, m, r):
145 async def _run(f, /, *a, r=r): return await r(f, *a)
146 cls.mgr, cls._run, cls.open_files = m, staticmethod(_run), {}
[docs]
147class MemoryMappedIOManager(A.LoopContextMixin):
148 __slots__ = '__factory', '__lock'
149 def __init__(self, executor=None, _f=(File,), _=H.create_executor): super().__init__(); self.__factory, self.__lock = type('__factory', _f, {}, m=__import__('_weakrefset').WeakSet(), r=partial(self.loop.run_in_executor, _(self, False) if executor is None else executor)), Lock()
150 @property
151 def open_maps(self): return self.__factory.mgr
152 def _run(self, f, /, *a): return self.__factory.run(f, *a)
153 @property
154 def currently_open(self): return len(self.open_maps)
155 @property
156 def open_paths(self): return dict(self.open_files.keys())
157 @property
158 def open_files(self): return self.__factory.open_files
159 @open_files.deleter
160 def open_files(self): self.open_files.clear()
161 @asynccontextmanager
162 async def _open(self, s, /, *a):
163 if (x := (F := self.open_files).get(a)): yield x; return
164 with await (r := self._run)(open, *a) as f:
165 if s > 0: await r(f.truncate, s)
166 async with self.__factory(f) as x:
167 F[a] = x
168 try: yield x
169 finally: F.pop(a, None)
[docs]
170 def open(self, path, init_size=0): return self._open(init_size, path, 'r+b')
[docs]
171 def create(self, path, init_size=0, *, exclusive=True): return self._open(init_size, path, 'x+b' if exclusive else 'w+b')
[docs]
172 async def __cleanup__(self):
173 async with self.__lock: self.open_maps.clear(); await gather(*(f.close() for f in self.open_files.values())); del self.open_files
[docs]
174 def __del__(self): self.make(self.__cleanup__())
[docs]
175 async def copy_file(self, src, dest, *, exclusive=True, flush=False):
176 async with self.open(src) as s, self.create(dest, exclusive=exclusive) as d:
177 await d.write(await s.read())
178 if flush: await d.flush()
[docs]
179 async def checksum(self, path, alg=None, _=partial(__import__('hashlib').new, usedforsecurity=False)):
180 async with self.open(path) as f: return await self._run(_, A.getcontext().MEMORY_MAPPED_IO_MANAGER_DEFAULT_CHECKSUM_ALG if alg is None else alg, await f.read()).hexdigest()
[docs]
181 async def approx_memory_usage(self):
182 async with self.__lock: return await self._run(self.__muh)
183 def __muh(self): return sum(m.size() for m in self.open_maps)
[docs]
184 @asynccontextmanager
185 async def prefetch_files(self, *P, init_size=0, _=S.exc_info):
186 l = tuple(map(partial(self.open, init_size=init_size), P))
187 try: yield await gather(*(c.__aenter__() for c in l)) # noqa: PLC2801
188 finally: t = _(); await gather(*(c.__aexit__(*t) for c in l))
[docs]
189 @asynccontextmanager
190 async def create_sparse_file(self, path, total_size, chunks):
191 async with self.create(path, total_size) as f:
192 g = f.smart_write
193 for o, d in chunks.items(): await g(d, o)
194 yield f
195 async def _bulk_reader(self, path, offsets):
196 a = (r := []).append
197 async with self.open(path) as f:
198 f = f.read
199 if offsets is None: a(await f())
200 else:
201 for o, s in offsets: a(await f(o, s))
202 return path, r
203 async def _bulk_writer(self, path, data):
204 async with self.open(path) as f: await gather(*starmap(f.write, data))
205 async def _bulk_creator(self, path, size, chunks, exclusive):
206 async with self.create(path, size, exclusive=exclusive) as f:
207 f = f.smart_write
208 for o, d in chunks.items(): await f(d, o)
209 async def _checksum_helper(self, alg, path): return path, await self.checksum(path, alg)
210 async def _resize_helper(self, path, size):
211 async with self.open(path) as f: await self._run(f.resize, size)
212 async def _compact_helper(self, path):
213 async with self.open(path) as f: await f.compact()
[docs]
214 async def bulk_read(self, file_offsets): return dict(await gather(*starmap(self._bulk_reader, file_offsets.items())))
[docs]
215 async def bulk_write(self, file_data): await gather(*starmap(self._bulk_writer, file_data.items()))
[docs]
216 async def bulk_checksum(self, paths, alg=None): return dict(await gather(*map(partial(self._checksum_helper, A.getcontext().MEMORY_MAPPED_IO_MANAGER_DEFAULT_CHECKSUM_ALG if alg is None else alg), paths)))
[docs]
217 async def bulk_copy(self, pairs): await gather(*starmap(self.copy_file, pairs))
[docs]
218 async def bulk_resize(self, sizes): await gather(*starmap(self._resize_helper, sizes.items()))
[docs]
219 async def compact_files(self, paths): await gather(*map(self._compact_helper, paths))
[docs]
220 async def find_in_files(self, pattern, paths, max_per_file=None, *, allow_overlapping=False):
221 async def search(p, o):
222 async with self.open(p) as f: return p, await (f.search if allow_overlapping else f.search_non_overlapping)(pattern, o, max_per_file)
223 return {k: v for k, v in await gather(*starmap(search, paths.items())) if v}
224 P.patch_method_signatures((__init__, 'executor=None'), (prefetch_files, '*paths, init_size=0'), (_open, 'init_size, path, mode, /'))
225def __getattr__(n, /, _=S, a=frozenset(('ainput', 'stdcoup')), g=globals()):
226 if n not in a: raise AttributeError(f'module {__name__!r} has no attribute {n!r}')
227 global ainput, stdcoup; stdcoup = AsyncReadWriteCouple(_.stdin, _.stdout) # ty: ignore[unresolved-global]
228 async def ainput(prompt='', assert_tty=False):
229 if prompt: await stdcoup.write(prompt); await stdcoup.flush()
230 if assert_tty and not stdcoup.reader.isatty(): raise OSError(E.ENOTTY, 'asyncutils.iotools.ainput: standard input is not a TTY')
231 if (d := await stdcoup.readline()).endswith('\n'): return d[:-1]
232 raise EOFError
233 return g[n]
234del f, H, P, S, z, File, O, t