如何编写恢复 cwd 的异步装饰器？

Question

我正在尝试为恢复cwd 的异步函数创建装饰器。我以此为参考：How do I write a decorator that restores the cwd?

这就是我想出的。它没有保留cwd，任何指针？

import os
import asyncio

def preserve_dir(function):
    async def decorator(*args, **kwargs):
        cwd = os.getcwd()
        result = function(*args, **kwargs)
        os.chdir(cwd)
        return await result
    return decorator


@preserve_dir
async def ordinary():
    os.chdir("/tmp")
    print(os.getcwd())


print(os.getcwd())
asyncio.run(ordinary())
print(os.getcwd())

Answer 1

当您的装饰器执行result = function(*args, **kwargs) 时，它不会得到结果。包装的function 只是返回一个协程对象——它实际上并没有运行函数体，直到你await 协程。通过在await 之前恢复工作目录，您将工作目录恢复得太早了。

不过，“在最后恢复这个全局状态”装饰器的整个概念是并发程序中的灾难。如果两个协程都使用您的装饰器并交错执行，则最终目录将取决于协程的调度方式。最终的目录可能不是你想要的，一个协程可能已经把工作目录从它下面改了出来。

不要使用装饰器，而是考虑编写独立于工作目录的代码，或者如果代码需要设置工作目录，则在工作进程中运行该代码，以便它拥有自己的工作目录。

Answer 2

正如其他人指出的那样，您需要在恢复工作目录之前等待修饰函数，因为调用异步函数不会执行它。

正如其他人也指出的那样，正确地执行此操作比看起来要困难得多，因为协程可以在运行时暂停到事件循环，并且在暂停时，不同的协程可以使用相同的装饰器更改目录。对于装饰器的简单实现，恢复原始协程会破坏它，因为工作目录将不再是它所期望的目录。理想情况下，您可以通过结构化代码来避免此问题，使其不依赖于当前工作目录。但从技术上讲， 可以实现一个正确的目录保留装饰器，只是需要额外的努力。虽然我不建议您在生产环境中执行此操作，但如果您对如何执行此操作感到好奇，请继续阅读。

This answer 展示了如何在每次恢复协程时应用上下文管理器。这个想法是创建一个协程包装器，一个等待它的__await__ 调用原始协程的__await__。通常这将使用yield from，但我们的包装器不会这样做，而是使用手写循环来模拟它，该循环使用send() 来恢复内部等待对象的迭代器。这提供了对内部等待的每次暂停和恢复的控制，用于在每次恢复时应用上下文管理器。请注意，这需要一个可重复使用的上下文管理器，它可以多次输入。

为了实现装饰器，我们需要一个可重用的目录保留上下文管理器，它不仅可以恢复__exit__ 中的前一个工作目录，还可以在下一个__enter__ 中重新应用它。前者在协程挂起（或返回）时恢复旧的工作目录，后者在协程恢复时恢复新的工作目录。装饰器只会将此上下文管理器传递给协程包装器：

# copy CoroWrapper from https://*.com/a/56079900/1600898

# context manager preserving the current directory
# can be re-entered multiple times
class PreserveDir:
    def __init__(self):
        self.inner_dir = None

    def __enter__(self):
        self.outer_dir = os.getcwd()
        if self.inner_dir is not None:
            os.chdir(self.inner_dir)

    def __exit__(self, *exc_info):
        self.inner_dir = os.getcwd()
        os.chdir(self.outer_dir)

def preserve_dir(fn):
    async def wrapped(*args, **kwds):
        return await CoroWrapper(fn(*args, **kwds), PreserveDir())
    return wrapped

此设置不仅通过了您的原始测试，还通过了一个更复杂的测试，该测试生成多个并发协同程序，这些协同程序使用相同的装饰器到不同的目录。例如：

@preserve_dir
async def ordinary1():
    os.chdir("/tmp")
    print('ordinary1', os.getcwd())
    await asyncio.sleep(1)
    print('ordinary1', os.getcwd())

@preserve_dir
async def ordinary2():
    os.chdir("/")
    print('ordinary2', os.getcwd())
    await asyncio.sleep(0.5)
    print('ordinary2', os.getcwd())
    await asyncio.sleep(0.5)
    print('ordinary2', os.getcwd())

async def main():
    await asyncio.gather(ordinary1(), ordinary2())

print(os.getcwd())
asyncio.run(main())
print(os.getcwd())

输出：

/home/user4815162342
ordinary1 /tmp
ordinary2 /
ordinary2 /
ordinary1 /tmp
ordinary2 /
/home/user4815162342

这种方法的一个警告是目录保存与当前的任务相关联。因此，如果您将执行委托给子协程，它会观察修改后的目录（如果它只是等待），但如果它使用 await asyncio.gather(coro1(), coro2()) 等待则不会。