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

问题描述

我正在尝试为还原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())

---

解决方案

正如其他人指出的那样，您需要等待修饰函数，然后才能还原工作目录，因为调用异步函数不会执行它。

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

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

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

# copy CoroWrapper from https://stackoverflow.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())等待，则不会观察修改后的目录。