通过函数调用使上下文管理器对象保持活动状态

Question

我在通过函数调用保持上下文管理器打开时遇到了一些问题。这就是我的意思：

在一个模块中定义了一个上下文管理器，我用它来打开与网络设备的 SSH 连接。 “设置”代码处理打开 SSH 会话并处理任何问题，而拆卸代码处理正常关闭 SSH 会话。我一般是这样使用的：

from manager import manager
def do_stuff(device):
    with manager(device) as conn:
        output = conn.send_command("show ip route")
        #process output...
    return processed_output 

为了保持 SSH 会话打开并且不必在函数调用之间重新建立它，我想向“do_stuff”添加一个参数，它可以选择返回 SSH 会话以及从 SSH 返回的数据会话，如下：

def do_stuff(device, return_handle=False):
    with manager(device) as conn:
        output = conn.send_command("show ip route")
        #process output...
        if return_handle:
            return (processed_output, conn)
        else:
            return processed_output

我希望能够从另一个函数调用此函数“do_stuff”，如下所示，这样它会向“do_stuff”发出信号，告知 SSH 句柄应与输出一起返回。

def do_more_stuff(device):
    data, conn = do_stuff(device, return_handle=True)
    output = conn.send_command("show users")
    #process output...
    return processed_output

但是我遇到的问题是 SSH 会话已关闭，这是由于 do_stuff 函数“返回”并触发了上下文管理器中的拆卸代码（这会优雅地关闭 SSH 会话）。

我尝试将“do_stuff”转换为生成器，使其状态暂停，并可能导致上下文管理器保持打开状态：

def do_stuff(device, return_handle=False):
    with manager(device) as conn:
        output = conn.send_command("show ip route")
        #process output...
        if return_handle:
            yield (processed_output, conn)
        else:
            yield processed_output

这样称呼它：

def do_more_stuff(device):
    gen = do_stuff(device, return_handle=True)
    data, conn = next(gen)
    output = conn.send_command("show users")
    #process output...
    return processed_output

但是这种方法在我的情况下似乎不起作用，因为上下文管理器被关闭，我得到一个关闭的套接字。

有没有更好的方法来解决这个问题？也许我的生成器需要更多的工作......我认为使用生成器来保持状态是我想到的最“明显”的方式，但总的来说，我是否应该研究另一种在函数调用之间保持会话打开的方式？

谢谢

Answer 1

我发现这个问题是因为我正在寻找一个类似问题的解决方案，其中我想要保持活动的对象是一个 pyvirtualdisplay.display.Display 实例，其中包含 selenium.webdriver.Firefox 实例。

如果在显示/浏览器实例创建期间引发异常，我还希望所有打开的资源都死掉。

我想这同样适用于您的数据库连接。

我认识到这可能只是部分解决方案，并且包含的​​做法并不理想。感谢您的帮助。

此答案是使用以下资源修补我的解决方案的广告库激增的结果：

• https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator
• http://www.wefearchange.org/2013/05/resource-management-in-python-33-or.html

（虽然我很欣赏它的潜力，但我还没有完全理解这里描述的内容。通过提供类似的情况，上面的第二个链接最终被证明是最有帮助的。）

from pyvirtualdisplay.display import Display
from selenium.webdriver import Firefox
from contextlib import contextmanager, ExitStack

RFBPORT = 5904


def acquire_desktop_display(rfbport=RFBPORT):
    display_kwargs = {'backend': 'xvnc', 'rfbport': rfbport}
    display = Display(**display_kwargs)
    return display


def release_desktop_display(self):
    print("Stopping the display.")
    # browsers apparently die with the display so no need to call quits on them
    self.display.stop()


def check_desktop_display_ok(desktop_display):
    print("Some checking going on here.")
    return True


class XvncDesktopManager:
    max_browser_count = 1

    def __init__(self, check_desktop_display_ok=None, **kwargs):
        self.rfbport = kwargs.get('rfbport', RFBPORT)
        self.acquire_desktop_display = acquire_desktop_display
        self.release_desktop_display = release_desktop_display

        self.check_desktop_display_ok = check_desktop_display_ok \
            if check_desktop_display_ok is None else check_desktop_display_ok

    @contextmanager
    def _cleanup_on_error(self):
        with ExitStack() as stack:
            """push adds a context manager’s __exit__() method
            to stack's callback stack."""
            stack.push(self)
            yield
            # The validation check passed and didn't raise an exception
            # Accordingly, we want to keep the resource, and pass it
            # back to our caller
            stack.pop_all()

    def __enter__(self):
        url = 'http://stackoverflow.com/questions/30905121/'\
            'keeping-context-manager-object-alive-through-function-calls'
        self.display = self.acquire_desktop_display(self.rfbport)
        with ExitStack() as stack:
            # add XvncDesktopManager instance's exit method to callback stack
            stack.push(self)
            self.display.start()

            self.browser_resources = [
                Firefox() for x in range(self.max_browser_count)
            ]

            for browser_resource in self.browser_resources:
                for url in (url, ):
                    browser_resource.get(url)

            """This is the last bit of magic.
            ExitStacks have a .close() method which unwinds
            all the registered context managers and callbacks
            and invokes their exit functionality."""
            # capture the function that calls all the exits
            # will be called later outside the context in which it was captured
            self.close_all = stack.pop_all().close
            # if something fails in this context in enter, cleanup
            with self._cleanup_on_error() as stack:
                if not self.check_desktop_display_ok(self):
                    msg = "Failed validation for {!r}"
                    raise RuntimeError(msg.format(self.display))
            # self is assigned to variable after "as",
            # manually call close_all to unwind callback stack
            return self

    def __exit__(self, *exc_details):
        # had to comment this out, unable to add this to callback stack
        # self.release_desktop_display(self)
        pass

我得到了一个半预期的结果：

    kwargs = {
        'rfbport': 5904,
    }
    _desktop_manager = XvncDesktopManager(check_desktop_display_ok=check_desktop_display_ok, **kwargs)
    with ExitStack() as stack:
        # context entered and what is inside the __enter__ method is executed
        # desktop_manager will have an attribute "close_all" that can be called explicitly to unwind the callback stack
        desktop_manager = stack.enter_context(_desktop_manager)


    # I was able to manipulate the browsers inside of the display
    # and outside of the context 
    # before calling desktop_manager.close_all()
    browser, = desktop_manager.browser_resources
    browser.get(url)
    # close everything down when finished with resource
    desktop_manager.close_all()  # does nothing, not in callback stack
    # this functioned as expected
    desktop_manager.release_desktop_display(desktop_manager)