# 每个资源都要自己管理，容易出错
f = None
try:
    f = open("data.txt")
    data = f.read()
    process(data)
finally:
    if f:
        f.close()

with open("data.txt") as f:
    data = f.read()
    process(data)
# 文件自动关闭

with EXPRESSION as TARGET:
    BLOCK

import dis


def demo():
    with open("/dev/null") as f:
        f.read()

dis.dis(demo)

# with EXPRESSION as TARGET:
#     BLOCK

# 等价于（简化版）：
manager = EXPRESSION
target = manager.__enter__()
exc = True
try:
    TARGET = target
    BLOCK
    exc = False
except:
    if not manager.__exit__(*sys.exc_info()):
        raise
finally:
    if not exc:
        manager.__exit__(None, None, None)

def __enter__(self):
    """进入 with 块时调用"""
    # 1. 执行准备工作（打开文件、获取锁、建立连接等）
    # 2. 返回一个值（可选），该值绑定到 as 子句
    return self  # 或其他对象

def __exit__(self, exc_type, exc_val, exc_tb):
    """离开 with 块时调用

    参数：
      exc_type: 异常类型（没有异常时为 None）
      exc_val:  异常实例（没有异常时为 None）
      exc_tb:   回溯对象（没有异常时为 None）

    返回值：
      True  → 抑制异常（异常不会继续传播）
      False → 让异常继续传播（默认行为）
    """
    # 1. 执行清理工作（关闭文件、释放锁、断开连接等）
    # 2. 根据异常情况决定返回值
    return False

class FileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()

        if exc_type is None:
            # 正常退出
            print("正常关闭")
            return False

        if exc_type is FileNotFoundError:
            # 文件未找到，记录日志但不中断
            print(f"文件不存在: {exc_val}")
            return True  # 抑制异常

        # 其他异常，记录并传播
        print(f"发生异常: {exc_type.__name__}: {exc_val}")
        return False  # 继续传播

class DatabaseConnection:
    """数据库连接上下文管理器"""

    def __init__(self, host, port, database):
        self.host = host
        self.port = port
        self.database = database
        self.connection = None

    def __enter__(self):
        print(f"连接到 {self.host}:{self.port}/{self.database}")
        # 模拟建立连接
        self.connection = {"host": self.host, "db": self.database}
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("关闭数据库连接")
        self.connection = None

        if exc_type:
            print(f"连接异常关闭: {exc_type.__name__}")
            # 数据库错误通常需要回滚，但这里继续传播
            return False

        print("连接正常关闭")
        return False


# 使用
with DatabaseConnection("localhost", 5432, "mydb") as conn:
    print(f"查询数据库: {conn}")

class ConnectionPool:
    """连接池 — 返回连接，而非自身"""

    def __init__(self):
        self._connections = [{"id": 1}, {"id": 2}]

    def __enter__(self):
        # 从池中借出一个连接
        conn = self._connections.pop()
        print(f"借出连接 {conn['id']}")
        return conn  # 返回连接对象，不是 self

    def __exit__(self, *args):
        # 归还连接
        self._connections.append(args)
        return False

from contextlib import contextmanager


@contextmanager
def managed_resource(*args, **kwargs):
    """用生成器实现上下文管理器"""
    # __enter__ 部分
    resource = acquire_resource(*args, **kwargs)
    try:
        yield resource  # ← 这个值被 as 绑定
    finally:
        # __exit__ 部分（无论如何都会执行）
        release_resource(resource)

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://python.org")) as page:
    for line in page:
        print(line)
# page.close() 自动调用

class closing:
    def __init__(self, thing):
        self.thing = thing
    def __enter__(self):
        return self.thing
    def __exit__(self, *exc_info):
        self.thing.close()

from contextlib import suppress

# 传统方式
try:
    os.remove("temp.txt")
except FileNotFoundError:
    pass

# suppress 方式
with suppress(FileNotFoundError):
    os.remove("temp.txt")

# 抑制多种异常
with suppress(FileNotFoundError, PermissionError):
    os.remove("protected.txt")

from contextlib import redirect_stdout, redirect_stderr
import io

# 捕获 print 输出到字符串
f = io.StringIO()
with redirect_stdout(f):
    print("这行不会显示在控制台")
    print("会被捕获到 StringIO")
captured = f.getvalue()  # "这行不会显示在控制台\n会被捕获到 StringIO\n"

# 重定向到文件
with open("log.txt", "w") as log, redirect_stdout(log):
    print("这行写入日志文件")

from contextlib import ExitStack


def process_files(file_list):
    """处理多个文件（数量在运行时才能确定）"""
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fname))
            for fname in file_list
        ]
        # 所有文件在退出时自动关闭
        for f in files:
            process(f)

# 1. 条件性进入
with ExitStack() as stack:
    if need_cache:
        cache = stack.enter_context(open("cache.txt"))
    # ... 其他操作

# 2. 推迟清理（手动清理）
stack = ExitStack()
f1 = stack.enter_context(open("a.txt"))
f2 = stack.enter_context(open("b.txt"))
# ... 一些操作
stack.close()  # 手动触发 LIFO 顺序清理

# 3. 注册回调
stack.callback(lambda: print("cleanup"))
stack.push(exit_func)  # 注册 __exit__ 风格的函数

# 4. pop_all 转移所有权
stack2 = ExitStack()
stack2.pop_all(stack)  # stack 的所有清理任务移交给 stack2

from contextlib import nullcontext


def process(data, profile=False):
    """profile=True 时启用性能分析，否则无操作"""
    profiler = cProfile.Profile() if profile else nullcontext()
    with profiler:
        result = expensive_computation(data)
    return result

from contextlib import ContextDecorator


class trace(ContextDecorator):
    def __enter__(self):
        print("开始")
        return self
    def __exit__(self, *exc):
        print("结束")
        return False


# 作为装饰器
@trace()
def hello():
    print("hello")

# 作为上下文管理器
with trace():
    print("world")

# 方式 1：缩进嵌套
with open("a.txt") as f1:
    with open("b.txt") as f2:
        data1 = f1.read()
        data2 = f2.read()

# 方式 2：单行多个（Python 3.1+）
with open("a.txt") as f1, open("b.txt") as f2:
    data1 = f1.read()
    data2 = f2.read()

# 方式 3：ExitStack（动态数量）
with ExitStack() as stack:
    f1 = stack.enter_context(open("a.txt"))
    f2 = stack.enter_context(open("b.txt"))

# 方式 4：括号语法（Python 3.10+）
with (
    open("a.txt") as f1,
    open("b.txt") as f2,
    open("c.txt") as f3,
):
    ...

class Tag:
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        print(f"<{self.name}>")
        return self
    def __exit__(self, *args):
        print(f"</{self.name}>")
        return False

with Tag("div"):
    with Tag("p"):
        print("Hello")
# 输出：
# <div>
# <p>
# Hello
# </p>
# </div>

class AsyncResource:
    async def __aenter__(self):
        await asyncio.sleep(0.1)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await asyncio.sleep(0.1)
        return False


# 使用
async def main():
    async with AsyncResource() as res:
        await res.do_work()

from contextlib import asynccontextmanager


@asynccontextmanager
async def async_open(filename):
    f = await aio_open(filename)
    try:
        yield f
    finally:
        await f.close()

class AtomicJSONWriter:
    """保证 JSON 写入的原子性

    原理：先写入临时文件，成功后再原子 rename
    如果中间出错，临时文件自动删除，不破坏原文件
    """

    def __init__(self, filepath, indent=2):
        self.filepath = filepath
        self.indent = indent
        self.temp_path = filepath + ".tmp"

    def __enter__(self):
        self.fp = open(self.temp_path, "w", encoding="utf-8")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fp.close()
        if exc_type is not None:
            os.remove(self.temp_path)  # 回滚
            return False
        os.replace(self.temp_path, self.filepath)  # 原子提交
        return False

    def write(self, data):
        json.dump(data, self.fp, ensure_ascii=False, indent=self.indent)

class ConnectionPool:
    """简单的数据库连接池管理器

    with pool as conn: 自动管理连接的获取和归还
    """

    def __init__(self, max_connections=5):
        self._pool = [self._create_conn(i) for i in range(max_connections)]
        self._in_use = set()
        self._lock = threading.Lock()

    def _create_conn(self, id):
        return {"id": id, "active": True}

    def __enter__(self):
        with self._lock:
            conn = self._pool.pop()
            self._in_use.add(conn["id"])
            return conn

    def __exit__(self, *args):
        with self._lock:
            self._in_use.remove(self._current_conn["id"])
            self._pool.append(self._current_conn)
        return False

class Timer:
    """既是上下文管理器又是装饰器的计时器"""

    def __init__(self, name="Timer"):
        self.name = name
        self.elapsed = None

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"[{self.name}] {self.elapsed*1000:.2f}ms")
        return False

    def __call__(self, func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return wrapper


# 两种用法
@Timer("排序")
def sort():
    ...

with Timer("查询"):
    ...

@contextmanager
def set_env(**environ):
    """临时设置环境变量"""
    old_environ = dict(os.environ)
    os.environ.update(environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(old_environ)


with set_env(DATABASE_URL="sqlite:///test.db"):
    assert os.environ["DATABASE_URL"] == "sqlite:///test.db"
# 环境变量恢复

class LogGroup:
    """日志分组 — 自动缩进和时间标记"""

    indent = 0

    def __init__(self, label):
        self.label = label
        self.start = None

    def __enter__(self):
        self.start = time.time()
        self._log(f"▶ {self.label}")
        LogGroup.indent += 1
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        LogGroup.indent -= 1
        elapsed = time.time() - self.start
        status = "✓" if exc_type is None else "✗"
        self._log(f"{status} {self.label} ({elapsed:.2f}s)")
        return False

    @classmethod
    def _log(cls, msg):
        print(f"{'  ' * cls.indent}{msg}")


with LogGroup("HTTP 请求"):
    LogGroup._log("验证 token...")
    with LogGroup("数据库查询"):
        LogGroup._log("SELECT ...")

class Resource:
    def __enter__(self):
        self._acquire()
        return self
    def __exit__(self, *args):
        self._release()
        return False

@contextmanager
def temporary_change(obj, attr, value):
    old = getattr(obj, attr)
    setattr(obj, attr, value)
    try:
        yield
    finally:
        setattr(obj, attr, old)

@contextmanager
def optional(condition, cm, else_cm=None):
    """条件性进入上下文管理器"""
    if condition:
        with cm:
            yield
    else:
        if else_cm:
            with else_cm:
                yield
        else:
            yield

@contextmanager
def transaction(conn):
    savepoint = conn.savepoint()
    try:
        yield
    except:
        conn.rollback(savepoint)
        raise
    else:
        conn.commit()

class Eater:
    def __exit__(self, *args):
        return True  # 吃掉一切异常

with Eater():
    1 / 0  # ZeroDivisionError 被静默吞掉
print("这行会执行")  # ← 真的会执行！

@contextmanager
def bad():
    print("enter")
    yield
    print("exit")  # 如果异常发生，这行不执行

@contextmanager
def good():
    print("enter")
    try:
        yield
    finally:
        print("exit")  # 无论如何都执行

# 方式 A
f = open("file")
try:
    process(f)
finally:
    f.close()

# 方式 B
with open("file") as f:
    process(f)

files = ["a.txt", "b.txt", "c.txt"]
# 想同时打开所有文件并处理

class Swallower:
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type == ValueError:
            print("吞掉 ValueError")
            return True
        # RuntimeError 不处理
        return False

with Swallower():
    raise ValueError("test")
print("ValueError 被吞掉")

with Swallower():
    pass  # 正常退出

# 装饰器方式 — 整个函数耗时
@timer
def load_data():
    ...

# 上下文管理器方式 — 代码块耗时
with timer():
    load_data()