如何借助第三方上下文管理器执行自定义类的方法？

问题背景 #

需要利用第三方包的上下文管理器以最佳实践处理资源清理，现有代码结构如下：

当前类实现 #

class MyClass(object):

    def __init__(self):
        self.open()

    def open(self):
        # 实例化对象 *例如* 打开文件

    def run(self):
        # 执行任务 *例如* 从文件句柄读取

    def cleanup(self):
        # 清理资源 *例如* 关闭文件句柄

现有调用方式 #

my_obj = MyClass()
my_obj.open()
my_obj.run()
my_obj.cleanup()

希望在类内部整合上下文管理器，让run()方法依赖上下文环境，已尝试两种方案：

方案1：外部显式调用类内的上下文管理器

my_obj2=MyClass2()
with my_obj2.third_party_context_manager_method() as mcmm:
    my_obj2.run()

方案2：上下文管理器封装在run方法内

my_obj3=MyClass3()
my_obj3.run_modified_containing_third_party_context_manager()

额外需求 #

实际场景中需要使用双重上下文管理器，该如何实现？

实际场景说明 #

实际为Dask集群/客户端的清理问题，需要确保Dask集群和客户端能被正确关闭，避免资源泄漏。

解决方案 #

1. 让自定义类实现上下文管理器协议 #

通过实现__enter__和__exit__方法，将第三方上下文管理器整合到类内部，用户只需通过with语句即可自动处理资源的进入和清理。

单个上下文管理器的实现示例

class MyClass:
    def __init__(self):
        self.third_party_ctx = None

    def __enter__(self):
        # 初始化并进入第三方上下文管理器
        self.third_party_ctx = some_third_party_context()  # 替换为实际第三方上下文
        self.third_party_ctx.__enter__()
        return self  # 返回类实例供外部调用run方法

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 退出第三方上下文管理器
        if self.third_party_ctx:
            self.third_party_ctx.__exit__(exc_type, exc_val, exc_tb)
        # 执行自定义清理逻辑
        self.cleanup()

    def run(self):
        # 直接依赖已初始化的上下文执行任务
        pass

    def cleanup(self):
        # 自定义资源清理操作
        pass

调用方式

with MyClass() as obj:
    obj.run()

2. 双重上下文管理器的实现 #

针对需要同时管理两个第三方上下文的场景（如Dask集群+客户端），有两种可靠实现方式：

方式一：手动按顺序管理上下文

from dask.distributed import LocalCluster, Client

class DaskManager:
    def __init__(self):
        self.cluster = None
        self.client = None

    def __enter__(self):
        # 先启动集群上下文
        self.cluster = LocalCluster()
        self.cluster.__enter__()
        # 再启动客户端上下文
        self.client = Client(self.cluster)
        self.client.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 按相反顺序退出上下文：先关闭客户端，再关闭集群
        if self.client:
            self.client.__exit__(exc_type, exc_val, exc_tb)
        if self.cluster:
            self.cluster.__exit__(exc_type, exc_val, exc_tb)
        self.cleanup()

    def run(self):
        # 执行Dask任务示例
        future = self.client.submit(lambda x: x + 1, 1)
        print(future.result())

    def cleanup(self):
        # 额外清理操作
        pass

方式二：使用ExitStack管理多上下文（推荐）

利用contextlib.ExitStack可以更灵活地管理任意数量的上下文，自动按相反顺序退出：

from contextlib import ExitStack
from dask.distributed import LocalCluster, Client

class DaskManager:
    def __init__(self):
        self.stack = ExitStack()
        self.cluster = None
        self.client = None

    def __enter__(self):
        # 将多个上下文加入栈，自动管理生命周期
        self.cluster = self.stack.enter_context(LocalCluster())
        self.client = self.stack.enter_context(Client(self.cluster))
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 关闭栈，自动退出所有上下文
        self.stack.close()
        self.cleanup()

    def run(self):
        future = self.client.submit(lambda x: x + 1, 1)
        print(future.result())

    def cleanup(self):
        pass

双重上下文的调用方式

with DaskManager() as dm:
    dm.run()

方案对比 #

• 方案1（外部显式调用上下文）：灵活性高，但需要用户手动处理上下文生命周期，容易遗漏清理操作。
• 方案2（上下文封装在run内）：调用简单，但每次调用run都会重新创建/销毁上下文，适合单次任务场景。
• 类实现上下文管理器：兼顾封装性和可靠性，用户只需通过with语句即可确保资源被正确清理，是最佳实践方案。

内容的提问来源于stack exchange，提问作者jtlz2