通过函数调用保持上下文管理器对象的存活

时间:2015-06-18 01:35:15

标签: python ssh network-programming paramiko contextmanager

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

我在模块中定义了一个上下文管理器,用于打开与网络设备的SSH连接。 “setup”代码处理打开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

然而,这种方法在我的情况下似乎不起作用,因为上下文管理器被关闭了,我回到了一个封闭的套接字。

有没有更好的方法来解决这个问题?也许我的生成器需要更多的工作...我认为使用生成器来保持状态是最明显的方式,但总的来说我应该考虑另一种方法来保持会话在函数调用中保持打开状态?

由于

1 个答案:

答案 0 :(得分:0)

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

如果在显示/浏览器实例创建期间引发异常,我还希望任何已打开的资源都会消失。

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

我认识到这可能只是一个部分解决方案并且包含不太好的做法。感谢帮助。

这个答案是ad lib峰值的结果,使用以下资源来修补我的解决方案:

(我还没有完全理解这里描述的内容,虽然我很欣赏它的潜力。上面的第二个链接最终证明是最有用的,提供了类似的情况。)

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)