03 · 装饰器与上下文管理器

# 装饰器就是 f = deco(f)：它在【定义那一刻】就执行了
def deco(fn):
    print("  [deco 被执行了] 拿到的是:", fn.__name__)
    def wrapper(*args, **kwargs):
        print("  [wrapper 被执行了]")
        return fn(*args, **kwargs)
    return wrapper

print("--- 开始定义函数 ---")

@deco
def add(a, b):
    return a + b

print("--- 定义完了，还没调用过 ---")
print("调用结果:", add(1, 2))

# 完全等价的手写形式
def sub(a, b):
    return a - b
sub = deco(sub)                 # ⭐ @deco 就是这一行的语法糖
print("手写等价:", sub(5, 3))

print("add 现在指向的是:", add.__name__)

# 「定义时执行」的用处和代价：注册表模式
REGISTRY = {}

def register(name):
    def decorator(fn):
        REGISTRY[name] = fn          # ⭐ 这一句在【定义那一刻】就跑了
        return fn                    # 原样返回，不包一层
    return decorator

@register("csv")
def load_csv(p):
    return f"csv:{p}"

@register("json")
def load_json(p):
    return f"json:{p}"

print("一次都没调用过，注册表已经满了:", list(REGISTRY))
print("按名字取:", REGISTRY["json"]("a.json"))
print("load_json 还是原函数吗:", load_json.__name__, "| 包了一层吗:", load_json is REGISTRY["json"])

# 不写 functools.wraps：函数的身份被整个吃掉
import functools
import inspect

def naked(fn):                          # ⚠️ 没有 wraps
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs)
    return wrapper

def wrapped(fn):                        # ⭐ 有 wraps
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs)
    return wrapper

def original(a, b=10, *, scale=1.0):
    """把 a 和 b 加起来再乘以 scale。"""
    return (a + b) * scale

A = naked(original)
B = wrapped(original)

for tag, f in [("原函数", original), ("naked", A), ("wraps", B)]:
    print(f"[{tag}]")
    print("   __name__  =", f.__name__)
    print("   __doc__   =", repr(f.__doc__))
    print("   签名       =", inspect.signature(f))
    print("   参数名     =", list(inspect.signature(f).parameters))

print("__qualname__ : naked ->", A.__qualname__, "| wraps ->", B.__qualname__)
print("__wrapped__  : naked 有吗", hasattr(A, "__wrapped__"),
      "| wraps 指回原函数", B.__wrapped__ is original)

# 带参数的装饰器：三层分别对应三次调用
import functools

def retry(times):                       # 第 1 层：收【装饰器的参数】
    print("  [1] retry(times) 执行，times =", times)
    def decorator(fn):                  # 第 2 层：收【被装饰的函数】
        print("  [2] decorator(fn) 执行，fn =", fn.__name__)
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):   # 第 3 层：收【调用时的实参】
            for i in range(1, times + 1):
                try:
                    return fn(*args, **kwargs)
                except ValueError as e:
                    print(f"  [3] 第 {i} 次失败: {e}")
            raise RuntimeError(f"{times} 次全失败")
        return wrapper
    return decorator

print("--- 定义 flaky ---")

calls = {"n": 0}

@retry(3)
def flaky():
    calls["n"] += 1
    if calls["n"] < 3:
        raise ValueError(f"第 {calls['n']} 次不行")
    return f"第 {calls['n']} 次成功"

print("--- 定义完了 ---")
print("结果:", flaky())

# ⭐ @retry(3) 展开就是这两步
def flaky2():
    return "ok"
flaky2 = retry(3)(flaky2)               # 先调 retry(3) 拿到 decorator，再调 decorator(fn)
print("手写等价:", flaky2())

# ⚠️ 少写一对括号：@retry 和 @retry(3) 差一层调用
import functools

def retry(times):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            for _ in range(times):
                return fn(*args, **kwargs)
        return wrapper
    return decorator

@retry                                  # ⚠️ 忘了 (3)，times 变成了函数本身
def ping():
    return "pong"

print("装饰这一步没报错，ping 现在是:", type(ping).__name__)
try:
    ping()
except TypeError as e:
    print("调用才炸:", type(e).__name__, ":", e)

# 想让两种写法都能用：判断第一个参数是不是可调用对象
def retry2(fn=None, *, times=3):
    if fn is None:                      # ⭐ 走 @retry2(times=5) 这条路
        return lambda f: retry2(f, times=times)
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        for _ in range(times):
            return fn(*args, **kwargs)
    return wrapper

@retry2
def a(): return "a 裸用"

@retry2(times=5)
def b(): return "b 带参数"

print(a(), "|", b())

# 叠几个装饰器：从下往上【应用】，从上往下【执行】
import functools

def mark(tag):
    def decorator(fn):
        print(f"  应用 {tag}，它拿到的是 {fn.__name__}")
        @functools.wraps(fn)
        def wrapper(*a, **k):
            print(f"  进入 {tag}")
            r = fn(*a, **k)
            print(f"  离开 {tag}")
            return r
        return wrapper
    return decorator

@mark("外层 A")
@mark("中层 B")
@mark("内层 C")
def core():
    print("  ★ 真正的函数体")
    return 42

print("--- 调用 ---")
print("返回:", core())

# 装饰器叠在类方法上：顺序错了会怎样
import functools

def log(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        print("    [log]", getattr(fn, "__name__", type(fn).__name__), "args=", args)
        return fn(*args, **kwargs)
    return wrapper

class C:
    @staticmethod                    # ⭐ 正确：staticmethod 在【最外层】
    @log
    def ok_static(x):
        return x * 2

    @log                             # ⚠️ 错误：log 拿到的是 staticmethod 对象
    @staticmethod
    def bad_static(x):
        return x * 2

    @classmethod                     # ⭐ 正确
    @log
    def ok_class(cls, x):
        return f"{cls.__name__}:{x}"

    @log                             # ⚠️ 错误
    @classmethod
    def bad_class(cls, x):
        return f"{cls.__name__}:{x}"

c = C()
print("ok_static :", C.ok_static(3))
print("ok_class  :", C.ok_class(3))

for name in ("bad_static", "bad_class"):
    try:
        print(f"{name}:", getattr(C, name)(3))
    except TypeError as e:
        print(f"{name}: TypeError :", e)

print("类字典里存的是什么类型:")
for name in ("ok_static", "bad_static", "ok_class", "bad_class"):
    print("  ", name, "->", type(C.__dict__[name]).__name__)

# bad_static 「从类上调能跑、从实例上调就废」
import functools
import sys

def log(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs)
    return wrapper

class C:
    @log
    @staticmethod                    # ⚠️ 顺序反了
    def bad_static(x):
        return x * 2

print("Python", sys.version.split()[0])
print("从类上调  C.bad_static(3) :", C.bad_static(3))
try:
    print("从实例上调 c.bad_static(3):", C().bad_static(3))
except TypeError as e:
    print("从实例上调 c.bad_static(3): TypeError :", e)

# with 背后就两个方法：__enter__ 和 __exit__
class Trace:
    def __init__(self, tag):
        self.tag = tag
    def __enter__(self):
        print(f"  [{self.tag}] __enter__ 跑了")
        return "⭐ as 绑的是 __enter__ 的【返回值】"   # ⚠️ 不是 self
    def __exit__(self, exc_type, exc_value, tb):
        print(f"  [{self.tag}] __exit__ 收到 exc_type =", exc_type)
        return False                                  # 不吞异常

with Trace("正常") as x:
    print("  body:", x)

print("---")
try:
    with Trace("出错") as x:
        raise ValueError("炸了")
except ValueError as e:
    print("异常照常传出来:", type(e).__name__, ":", e)

print("---")
def f():
    with Trace("提前 return"):
        return "返回值"                                # ⭐ __exit__ 照样跑
print("f() =", f())

print("---")
for i in range(3):
    with Trace(f"break-{i}"):
        if i == 1:
            break                                      # ⭐ __exit__ 照样跑

# with 就是 try/finally 的封装版；多个资源时 with 比手写可靠
from contextlib import ExitStack

class Res:
    def __init__(self, name, fail=False):
        self.name, self.fail = name, fail
    def __enter__(self):
        if self.fail:
            raise OSError(f"{self.name} 打不开")
        print("  开", self.name)
        return self
    def __exit__(self, *exc):
        print("  关", self.name)
        return False

print("=== 手写 try/finally（等价物）===")
r = Res("A")
x = r.__enter__()                       # ⭐ with 语句展开后就是这三段
try:
    print("  用", x.name)
finally:
    r.__exit__(None, None, None)

print("=== 一行开两个：B 成功、C 失败 ===")
try:
    with Res("B") as b, Res("C", fail=True) as c:
        print("  到不了这里")
except OSError as e:
    print("  捕获:", e, "—— ⭐ 注意 B 已经被关掉了")

print("=== 数量不定：ExitStack ===")
names = ["D", "E", "F"]
with ExitStack() as stack:
    opened = [stack.enter_context(Res(n)) for n in names]
    print("  一共开了", len(opened), "个")
print("  ⭐ 退出时按【相反顺序】全部关掉")

# 💀 __exit__ 返回 True 会把异常吞掉，而且不留任何痕迹
import contextlib

class Swallow:
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, tb):
        return True                 # ⚠️ 真值 =「这个异常我处理掉了」

class Propagate:
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, tb):
        return None                 # ⭐ None / False = 让异常继续往外走

def save(cm, tag):
    written = []
    try:
        with cm:
            written.append("第一条")
            raise OSError("磁盘满了")
            written.append("第二条")       # 永远到不了
        print(f"{tag}: ⚠️ with 块【正常结束】了，代码继续往下跑")
    except OSError as e:
        print(f"{tag}: 外面捕获到 {type(e).__name__}: {e}")
    print(f"{tag}: 实际写进去 {len(written)} 条 ->", written)

save(Swallow(), "Swallow  ")
print("---")
save(Propagate(), "Propagate")

print("--- 真想吞，用 contextlib.suppress，写明吞哪一种 ---")
with contextlib.suppress(FileNotFoundError):
    open("绝对不存在的文件.txt")
print("FileNotFoundError 被吞了，程序还在跑")
try:
    with contextlib.suppress(FileNotFoundError):
        raise PermissionError("别的异常照样往外走")
except PermissionError as e:
    print("PermissionError 没被吞:", e)

# contextlib.contextmanager：yield 把函数劈成「进入前」和「退出后」
from contextlib import contextmanager

@contextmanager
def naive(tag):
    print(f"  [{tag}] 进入前")
    yield tag                        # ⚠️ 没有 try/finally
    print(f"  [{tag}] 退出后 —— 只有【没出异常】时才跑得到")

@contextmanager
def robust(tag):
    print(f"  [{tag}] 进入前")
    try:
        yield tag                    # ⭐ body 里的异常是从这一行【抛回来】的
    except ValueError as e:
        print(f"  [{tag}] 在 yield 处接住了 {type(e).__name__}: {e}")
        raise                        # 不 raise 就等于吞掉，和 __exit__ 返回 True 一样
    finally:
        print(f"  [{tag}] finally 清理 —— 出不出异常都跑")

print("=== naive，正常路径 ===")
with naive("naive-ok") as t:
    print("  body:", t)

print("=== naive，body 抛异常 ===")
try:
    with naive("naive-bad"):
        raise ValueError("炸了")
except ValueError:
    print("  外面接住了 —— ⚠️ 但「退出后」那行【从来没跑过】")

print("=== robust，body 抛异常 ===")
try:
    with robust("robust"):
        raise ValueError("炸了")
except ValueError:
    print("  外面接住了")

# contextmanager 的三个坑：yield 前的异常 / 两个 yield / 一次性
from contextlib import contextmanager

@contextmanager
def check(path):
    if not path.endswith(".csv"):
        raise ValueError("只支持 csv")     # ⭐ 在 __enter__ 里抛，body 根本不会开始
    print("  打开", path)
    try:
        yield path
    finally:
        print("  关闭", path)

try:
    with check("data.txt"):
        print("  ⚠️ 这一行不会出现")
except ValueError as e:
    print("① with 那一行就抛了:", type(e).__name__, ":", e)

@contextmanager
def two_yields():
    yield 1
    yield 2                                # ⚠️ 只允许一个 yield

try:
    with two_yields():
        pass
except RuntimeError as e:
    print("② 两个 yield ->", type(e).__name__, ":", e)

# ⚠️ contextmanager 造出来的对象是【一次性】的（它底下就是个生成器）
cm = check("a.csv")
with cm:
    pass
try:
    with cm:                               # 第二次
        pass
except Exception as e:
    print("③ 同一个对象用第二次 ->", type(e).__name__)

# ⭐ 同一个东西既能 with 也能 @：contextmanager 造出来的对象自带 ContextDecorator
from contextlib import contextmanager
import time

@contextmanager
def timer(tag):
    t0 = time.perf_counter()
    try:
        yield
    finally:
        print(f"  [{tag}] 耗时 {(time.perf_counter() - t0) * 1000:.1f} ms")

with timer("当 with 用"):
    sum(range(1_000_000))

@timer("当装饰器用")                  # ⭐ 同一个对象，直接贴到函数上
def work():
    return sum(range(1_000_000))

work()
work()                                # ⭐ 能用第二次：每次调用内部会【重建】一个生成器

print("timer('x') 是什么:", type(timer("x")).__name__)

# 🧩 骨架：`bn` 来自你自己的代码，这一段只看写法
bn.train()                                   # ⚠️ 忘了切回 eval()
with torch.no_grad():                        # ⚠️ 而且包在 no_grad 里，看起来很安全
    bn(val_x)
bn.eval()

# 成对操作写成手工的「设置…还原」，中途一炸就还原不了
from contextlib import contextmanager

class FakeModel:
    def __init__(self):
        self.training = True
    def eval(self):
        self.training = False
    def train(self):
        self.training = True
    def forward(self, x):
        if x < 0:
            raise ValueError("坏样本")
        return x * 2

m = FakeModel()

def validate_manual(model, batch):
    model.eval()
    out = [model.forward(x) for x in batch]
    model.train()                       # ⚠️ 抛异常就跑不到这一行
    return out

try:
    validate_manual(m, [1, 2, -1])
except ValueError:
    pass
print("手工还原：出异常后 model.training =", m.training, "⚠️ 卡在 eval 模式了")

@contextmanager
def eval_mode(model):
    was = model.training
    model.eval()
    try:
        yield model
    finally:
        model.training = was            # ⭐ 出不出异常都还原

m.train()
def validate_with(model, batch):
    with eval_mode(model):
        return [model.forward(x) for x in batch]

try:
    validate_with(m, [1, 2, -1])
except ValueError:
    pass
print("with 还原：出异常后 model.training =", m.training, "⭐ 还原了")