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

2022-03-25 00:00:00 python python-3.x python-asyncio

问题描述

我正在尝试为还原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())等待,则不会观察修改后的目录。

相关文章