Python 的上下文管理

def handle_request(trace_id: str) -> None:
    load_user(trace_id)
    create_order(trace_id)

from contextvars import ContextVar

trace_id: ContextVar[str] = ContextVar("trace_id")


def write_log(message: str) -> None:
    print(f"[{trace_id.get('unknown')}] {message}")


def load_order() -> None:
    write_log("load order")


def handle_request(incoming_trace_id: str) -> None:
    token = trace_id.set(incoming_trace_id)
    try:
        write_log("request started")
        load_order()
    finally:
        trace_id.reset(token)


handle_request("4bf92f3577b34da6")
print(f"outside: {trace_id.get('MISSING')}")

[4bf92f3577b34da6] request started
[4bf92f3577b34da6] load order
outside: MISSING

from contextvars import Context, ContextVar, copy_context
from threading import Thread

trace_id: ContextVar[str] = ContextVar("trace_id", default="MISS")


def show_trace(name: str) -> None:
    print(f"{name}: {trace_id.get()}")


trace_id.set("req-42")

# 场景一：复制主线程的当前 Context，让新线程继承 trace_id。
copied_context = copy_context()
copied = Thread(
    target=copied_context.run,  # target：新线程启动后调用的函数
    args=(show_trace, "copied"),  # args：传给 Context.run 的位置参数
)
copied.start()
copied.join()

# 场景二：从空 Context 启动，新线程看不到主线程的 trace_id。
empty_context = Context()
isolated = Thread(
    target=empty_context.run,
    args=(show_trace, "isolated"),
)
isolated.start()
isolated.join()

copied: req-42  # 复制的 Context 已包含主线程的 trace_id
isolated: MISS  # 空 Context 没有绑定，返回 ContextVar 的默认值

Thread(target=show_trace, args=("copied",), context=copy_context())
Thread(target=show_trace, args=("isolated",), context=Context())

from concurrent.futures import ThreadPoolExecutor
from contextvars import copy_context
from typing import Any, Callable


def submit_with_context(
    executor: ThreadPoolExecutor,
    func: Callable[..., Any],
    /,
    *args: Any,
    **kwargs: Any,
):
    context = copy_context()
    return executor.submit(context.run, func, *args, **kwargs)

import asyncio
from contextvars import ContextVar

value: ContextVar[str] = ContextVar("value", default="MISS")


async def direct() -> None:
    print(f"direct before: {value.get()}")
    value.set("direct")  # 故意不 reset，用来观察与调用方共享的绑定
    print(f"direct after: {value.get()}")


async def child(start: asyncio.Event) -> None:
    await start.wait()  # 等父 Task 写入 parent-2 后再读取
    print(f"child before: {value.get()}")
    value.set("child")
    print(f"child after: {value.get()}")


async def main() -> None:
    # 场景一：直接 await 不创建新 Task，direct() 沿用当前 Context。
    value.set("parent")
    await direct()
    print(f"parent after direct: {value.get()}")

    # 场景二：create_task() 在此处复制当前 Context，其中 value 为 direct。
    start = asyncio.Event()
    task = asyncio.create_task(child(start))
    value.set("parent-2")  # 只修改父 Task 的 Context
    start.set()  # 放行子 Task，让它读取创建时捕获的 Context
    await task
    print(f"parent after task: {value.get()}")


asyncio.run(main())

direct before: parent       # direct() 读取父 Task 当前的绑定
direct after: direct        # direct() 在同一个 Context 中写入 direct
parent after direct: direct # 直接 await 没有复制，父协程读到新值
child before: direct        # 子 Task 读取创建时捕获的 direct
child after: child          # 子 Task 只修改自己的 Context
parent after task: parent-2 # 子 Task 的写入没有回传父 Task

from concurrent.futures import ProcessPoolExecutor
from contextvars import ContextVar
from multiprocessing import get_context

trace_id: ContextVar[str] = ContextVar("trace_id", default="MISS")


def read_implicit() -> str:
    # spawn worker 重新导入模块，本地 Context 中没有父进程的绑定。
    return trace_id.get()


def read_explicit(incoming_trace_id: str) -> str:
    # 用普通任务参数在 worker 的 Context 中重建绑定。
    token = trace_id.set(incoming_trace_id)
    try:
        return trace_id.get()
    finally:
        trace_id.reset(token)


def main() -> None:
    trace_id.set("req-42")
    with ProcessPoolExecutor(
        max_workers=1,
        mp_context=get_context("spawn"),  # 启动全新的 Python 解释器
    ) as executor:
        # 场景一：任务没有接收 trace_id，只能读取 worker 的默认值。
        implicit = executor.submit(read_implicit).result()

        # 场景二：把 trace_id 作为可序列化参数传入 worker。
        explicit = executor.submit(read_explicit, trace_id.get()).result()

    print(f"implicit: {implicit}")
    print(f"explicit: {explicit}")


if __name__ == "__main__":
    main()

implicit: MISS   # 没有传播父进程 Context，返回默认值
explicit: req-42 # worker 用任务参数重建了本地绑定

def get(var, argument_default=NO_ARGUMENT):
    tstate = current_thread_state()
    context = tstate.context

    # 普通 GIL 构建先用线程状态 ID 和 Context 版本校验缓存。
    if cache_matches(var, tstate):
        return var.cached_value

    if context is not None and var in context.ctx_vars:
        value = context.ctx_vars[var]
        refresh_cache(var, tstate, value)
        return value

    # 两级默认值只参与本次读取，不会写入 Context。
    if argument_default is not NO_ARGUMENT:
        return argument_default
    if var.has_default:
        return var.default

    return NOT_FOUND

def set(var, value):
    context = context_get()  # 没有当前 Context 时，在此处延迟创建。
    old = context.ctx_vars.get(var, MISSING)
    token = Token(context, var, old)

    # assoc() 返回新的 HAMT，未变化的节点仍与旧映射共享。
    context.ctx_vars = context.ctx_vars.assoc(var, value)
    return token


def reset(var, token):
    context = context_get()
    validate_token(token, var, context)
    token.used = True

    if token.old is MISSING:
        # set() 前没有绑定，reset() 应删除键，而不是写入默认值。
        context.ctx_vars = context.ctx_vars.without(var)
    else:
        context.ctx_vars = context.ctx_vars.assoc(var, token.old)


def copy_context():
    context = context_get()

    # 新 Context 只共享不可变 HAMT，不逐项复制其中的绑定。
    return new_context(shared_vars=context.ctx_vars)

def enter(ctx):
    tstate = current_thread_state()
    if ctx.ctx_entered:
        raise RuntimeError("Context 已经进入")

    # 保存旧 Context，再把 ctx 压入当前线程的 Context 栈。
    ctx.ctx_prev = tstate.context
    tstate.context = ctx
    ctx.ctx_entered = True

    # 版本变化会使 ContextVar 上的旧读取缓存失效。
    tstate.context_ver += 1


def exit(ctx):
    tstate = current_thread_state()
    if not ctx.ctx_entered:
        raise RuntimeError("Context 尚未进入")

    # 弹出 ctx，恢复进入前保存的 Context。
    tstate.context = ctx.ctx_prev
    ctx.ctx_prev = None
    ctx.ctx_entered = False
    tstate.context_ver += 1

def create_task(coro, context=None):
    task = allocate_task(coro)

    # 未显式传 context 时，只在 Task 创建阶段复制一次。
    task.task_context = (
        copy_current_context()
        if context is None
        else context
    )

    # Handle 保存 task_context，等事件循环调度时再进入它。
    loop.call_soon(task_step, task, context=task.task_context)
    return task


def run_handle(handle):
    # 标准 Handle._run() 通过 Context.run() 恢复注册时的 Context。
    handle.context.run(handle.callback, *handle.args)


def task_step(task, error=None):
    if error is None:
        # C 加速路径在这里调用 PyIter_Send(task.coro, None)。
        task.coro.send(None)
    else:
        # Future 的异常或取消会通过 throw() 注入协程。
        task.coro.throw(error)

def wait_for_future(task, future):
    # Task 挂起时，把唤醒回调和同一个 task_context 一起注册。
    future.add_done_callback(
        task.wakeup,
        context=task.task_context,
    )


def schedule_future_callbacks(future):
    for callback, context in future.callbacks:
        # Future 完成后，Context 会随回调进入事件循环队列。
        loop.call_soon(callback, future, context=context)


def task_wakeup(task, future):
    # Handle 已进入 task_context，此处只需检查结果并恢复 Task。
    try:
        future.result()
    except BaseException as error:
        task_step(task, error)
    else:
        task_step(task, None)

from contextvars import ContextVar
from typing import Awaitable, Callable, TypeVar

Response = TypeVar("Response")
trace_id: ContextVar[str] = ContextVar("trace_id")


async def run_request(
    incoming_trace_id: str,
    handler: Callable[[], Awaitable[Response]],
) -> Response:
    token = trace_id.set(incoming_trace_id)
    try:
        return await handler()
    finally:
        trace_id.reset(token)