python - 异步子进程上下文管理器

Question

我想要一个异步上下文管理器，以确保托管进程在通过正常或强制终止退出之前完成，而不管取消如何。例如，

async def sleep(duration):
  proc = await asyncio.create_subprocess_exec('sleep', str(duration))
  async with manage_process(proc):
     ...
  # proc either finished, was terminated, or killed.

到目前为止，我得到的最好的结果是：

@asynccontextmanager
async def manage_process(
    proc: asyncio.subprocess.Process,
    timeout: Optional[float] = None,
) -> AsyncIterator[asyncio.subprocess.Process]:
    if timeout is None:
        timeout = 1.0

    try:
        yield proc
    finally:
        try:
            proc.terminate()
        except ProcessLookupError:
            pass
        else:
            try:
                await asyncio.shield(asyncio.wait_for(proc.wait(), timeout))
            except asyncio.CancelledError:
                is_done = False
                is_cancelled = True
            except asyncio.TimeoutError:
                is_done = False
                is_cancelled = False
            else:
                is_done = True
                is_cancelled = False

            if not is_done:
                try:
                    proc.kill()
                except ProcessLookupError:
                    pass
                else:
                    await asyncio.shield(proc.wait())

            if is_cancelled:
                raise asyncio.CancelledError()

我不确定我是否正确处理了取消。我没用过asyncio.shield()，我不确定我做得对。

标签： pythonsubprocesspython-asynciocontextmanager