如何禁用嵌套上下文管理器？自定义cast与no_cast实现咨询

要实现你需要的no_cast上下文管理器，核心思路是通过全局状态标记让cast在执行生效逻辑前先检查是否处于no_cast模式，若处于该模式则跳过自身的类型转换配置。以下是具体实现步骤：

1. 新增全局状态管理 #

首先添加模块级别的全局变量和辅助函数，用于跟踪no_cast和原有cast的状态：

from typing import Any
import numpy as np

# 全局状态：标记是否处于no_cast模式
_in_no_cast = False
# 原有的cast状态管理
_cast_enabled = False
_cast_dtype = None

def is_cast_enabled() -> bool:
    return _cast_enabled

def set_cast_enabled(enabled: bool) -> None:
    global _cast_enabled
    _cast_enabled = enabled

def set_cast_dtype(dtype) -> None:
    global _cast_dtype
    _cast_dtype = dtype

# 新增的no_cast状态管理
def is_no_cast_enabled() -> bool:
    return _in_no_cast

def set_no_cast_enabled(enabled: bool) -> None:
    global _in_no_cast
    _in_no_cast = enabled

2. 实现no_cast上下文管理器 #

这个管理器负责在进入时开启no_cast模式，退出时恢复之前的状态（支持多层嵌套使用）：

class no_cast:
    def __enter__(self) -> None:
        # 保存进入前的no_cast状态
        self.prev_no_cast = is_no_cast_enabled()
        set_no_cast_enabled(True)
    
    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        # 恢复之前的no_cast状态
        set_no_cast_enabled(self.prev_no_cast)

3. 修改cast上下文管理器 #

在cast的__enter__方法中加入no_cast模式检查，若处于该模式则跳过自身的生效逻辑：

class cast:
    def __init__(self, enabled: bool = True, dtype) -> None:
        self.prev = False
        self.enabled = enabled
        self.dtype = dtype

    def __enter__(self) -> None:
        self.prev = is_cast_enabled()
        # 仅当不在no_cast模式时，才执行cast的配置逻辑
        if not is_no_cast_enabled():
            set_cast_enabled(self.enabled)
            set_cast_dtype(self.dtype)

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        # 退出时恢复之前的cast状态，确保状态不混乱
        set_cast_enabled(self.prev)

4. 测试验证 #

用示例代码验证效果：

# 模拟测试数据
a = np.array([1.0, 2.0], dtype=np.float32)
b = np.array([3.0, 4.0], dtype=np.float32)

# 正常cast场景：结果转为float16
with cast(dtype=np.float16):
    c = a + b
print(f"正常cast结果类型：{c.dtype}")  # 输出 float16

# no_cast嵌套cast场景：结果保持原float32类型
with no_cast():
    with cast(dtype=np.float16, enabled=True):
        c = a + b
print(f"no_cast内cast结果类型：{c.dtype}")  # 输出 float32

原理说明 #

• no_cast通过全局标记_in_no_cast控制所有嵌套cast的行为，只要该标记为True，cast就会跳过自身的类型转换配置，无论其enabled参数设为True还是False。
• 上下文管理器的__exit__方法负责恢复进入前的状态，确保嵌套使用或异常场景下全局状态不会混乱。

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