Python上下文管理器：如何优雅关联两个对象？自定义类实现需求

问题描述 #

在阅读TensorFlow文档时看到这段代码：

g = tf.Graph()
with g.as_default():
  c = tf.constant(5.0)
  assert c.graph is g

tf.Graph().as_default()返回一个上下文管理器，with块内定义的对象会自动关联到g，无需显式传入。想理解这种无需显式传参的对象关联逻辑，并实现两个类A和B，满足以下要求：

• A作为上下文管理器
• 在with A()块内创建B实例时，B的a属性自动关联到当前的A实例
• 目标运行代码：

obj = A()
with obj as a:
    b = B()
    assert b.a is obj

给定的类框架：

class A:
    def __init__(self):
        pass

    def __enter__(self):
        # something here

    def __exit__(self):
        # something here

class B:
    def __init__(self):
        self.a = None

实现原理与代码 #

这种自动关联的核心是线程局部存储（Thread-Local Storage）——它能在当前线程中保存一个“全局”状态，但这个状态只对当前线程可见，不会和其他线程冲突。TensorFlow就是用类似机制记录当前默认图的。

完整实现代码 #

import threading

class A:
    # 线程局部存储，用于保存当前活跃的A实例
    _current_a = threading.local()

    def __init__(self):
        pass

    def __enter__(self):
        # 保存之前的活跃实例（处理嵌套上下文场景）
        self._prev_a = getattr(A._current_a, 'instance', None)
        # 将当前A实例设为线程内的活跃实例
        A._current_a.instance = self
        # 返回当前实例（对应with语句里的as a）
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 恢复之前的活跃实例，避免影响后续代码
        A._current_a.instance = self._prev_a

class B:
    def __init__(self):
        # 从线程局部存储中获取当前活跃的A实例
        self.a = getattr(A._current_a, 'instance', None)

代码说明 #

1. 线程局部存储_current_a：A类定义了类级别的threading.local对象，用来存储当前线程中正在使用的A实例，每个线程的存储区独立，不会互相干扰。
2. __enter__方法：进入上下文时，先保存之前的活跃实例（支持嵌套with场景），再将当前A实例设为线程内的活跃实例，最后返回自身供with ... as a绑定使用。
3. __exit__方法：退出上下文时，恢复之前的活跃实例，确保上下文结束后不会污染后续的线程状态。
4. B的__init__：创建B实例时，直接从A的线程局部存储中读取当前活跃的A实例，赋值给self.a，实现自动关联。

测试验证 #

运行目标代码：

obj = A()
with obj as a:
    b = B()
    assert b.a is obj
    print("断言成功：b.a 与 obj 是同一个实例")

# 上下文外创建B，a属性为None
b_out = B()
print(f"上下文外的B实例a属性：{b_out.a}")

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