async def rmtree(self,
sema: Optional[asyncio.Semaphore],
url: str,
listener: Optional[Callable[[int], None]] = None) -> None:
if listener is None:
listener = lambda _: None # noqa: E731
if sema is None:
sema = asyncio.Semaphore(50)
async def rm(entry: FileListEntry):
assert listener is not None
listener(1)
await self._remove_doesnt_exist_ok(await entry.url())
listener(-1)
try:
it = await self.listfiles(url,
recursive=True,
exclude_trailing_slash_files=False)
except FileNotFoundError:
return
async with OnlineBoundedGather2(sema) as pool:
tasks = [pool.call(rm, entry) async for entry in it]
if tasks:
await pool.wait(tasks)
async def rmtree(self, sema: asyncio.Semaphore, url: str) -> None:
async with OnlineBoundedGather2(sema) as pool:
try:
it = await self.listfiles(url, recursive=True)
except FileNotFoundError:
return
async for entry in it:
await pool.call(self._remove_doesnt_exist_ok, await entry.url())
async def _rmtree(self, sema: asyncio.Semaphore, url: str) -> None:
async with OnlineBoundedGather2(sema) as pool:
bucket, name = self._get_bucket_name(url)
if name and not name.endswith('/'):
name = f'{name}/'
it = self._listfiles_recursive(bucket, name)
async for entry in it:
await pool.call(self._remove_doesnt_exist_ok, await
entry.url())
async def __aexit__(self, exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType]) -> None:
async with OnlineBoundedGather2(self._sema) as pool:
try:
if exc_val is not None:
return
async def tree_compose(names, dest_name):
n = len(names)
assert n > 0
if n <= 32:
await self._compose(names, dest_name)
return
q, r = divmod(n, 32)
i = 0
p = 0
chunks = []
while i < 32:
# each chunk gets q, and the first r get one more
chunk_size = q
if i < r:
chunk_size += 1
chunks.append(names[p:p + chunk_size])
p += chunk_size
i += 1
assert p == n
assert len(chunks) == 32
chunk_names = [
self._tmp_name(f'chunk-{secret_alnum_string()}')
for _ in range(32)
]
chunk_tasks = [
await pool.call(tree_compose, c, n)
for c, n in zip(chunks, chunk_names)
]
await pool.wait(chunk_tasks)
await self._compose(chunk_names, dest_name)
for n in chunk_names:
await pool.call(self._fs._remove_doesnt_exist_ok,
f'gs://{self._bucket}/{n}')
await tree_compose(
[self._part_name(i) for i in range(self._num_parts)],
self._dest_name)
finally:
await self._fs.rmtree(
self._sema,
f'gs://{self._bucket}/{self._dest_dirname}_/{self._token}')
async def __aexit__(self, exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType]) -> None:
async with OnlineBoundedGather2(self._sema) as pool:
cleanup_tasks = []
try:
if exc_val is not None:
return
async def tree_compose(names, dest_name):
n = len(names)
assert n > 0
if n <= 32:
await self._compose(names, dest_name)
return
q, r = divmod(n, 32)
i = 0
p = 0
chunks = []
while i < 32:
# each chunk gets q, and the first r get one more
chunk_size = q
if i < r:
chunk_size += 1
chunks.append(names[p:p + chunk_size])
p += chunk_size
i += 1
assert p == n
assert len(chunks) == 32
chunk_names = [
self._tmp_name(f'chunk-{secret_alnum_string()}')
for _ in range(32)
]
chunk_tasks = [
await pool.call(tree_compose, c, n)
for c, n in zip(chunks, chunk_names)
]
await pool.wait(chunk_tasks)
await self._compose(chunk_names, dest_name)
for n in chunk_names:
cleanup_tasks.append(await pool.call(
self._fs._remove_doesnt_exist_ok,
f'gs://{self._bucket}/{n}'))
await tree_compose(
[self._part_name(i) for i in range(self._num_parts)],
self._dest_name)
finally:
await self._fs.rmtree(
self._sema, f'gs://{self._bucket}/{self._dest_dirname}_')
# after the rmtree, all temporary files should be gone
# cancel any cleanup tasks that are still running
# exiting the pool will wait for everything to finish
for t in cleanup_tasks:
if not t.done():
t.cancel()