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()