使用上下文管理器选择数据库或文件处理程序

Question

我希望能够使用上下文管理器根据参数打开 FileHandler 或 DBHandler。 2 类本身就是上下文管理器。我可以想出下面的代码，想知道是否有更好的方法来做到这一点（请忽略任何遗漏的功能）？

class FileHandler:
    def __init__(self, name):
        self._file = open(name, 'w+')
    def write(self, val):
        self._file.write(val)
    def close(self):
        self._file.close()
    def __enter__(self):
        return self
    def __exit__(self, *exc):
        return False

class DBHandler:
    def __init__(self, name):
        self._db = some_db_api.open(name)
    def write(self, val):
        # val is some query
        self._db.execute(val)
    def close(self):
        self._db.commit()
    def __enter__(self):
        return self
    def __exit__(self, *exc):
        return False
        
from contextlib import contextmanager

@contextmanager
def get_resource(resource_type, *args, **kwargs):
    try:
        if resource_type == 'file':
            resource = FileHandler(*args, **kwargs)
        else:
            resource = DBHandler(*args, **kwargs)
        yield resource
        resource.close()
    except:
        resource.rollback()


if output_mode == 'f':
    with get_resource('file', name) as resource:
        for i in range(5):
            resource.write(str(i))
else:
    with get_resource('db', name) as resource:
        for i in range(5):
            resource.write(str(i))

Answer 1

with 语句不需要是上下文管理器类型的可调用对象——任何 计算为 的表达式上下文管理器工作。这允许函数调用、条件表达式以及直接名称引用来提供上下文管理器。

您几乎可以删除所有间接寻址，并直接 select 适当的上下文管理器：

class DBHandler:
    def __init__(self, name):
        self._db = some_db_api.open(name)
    def write(self, val):
        self._db.execute(val)
    def __enter__(self):
        return self
    def __exit__(self, *exc):
        # make a commit of no error occurred, rollback otherwise
        if exc[0] is None:
            self._db.commit()
        else:
            self._db.rollback()

context = open(name) if output_mode == 'f' else DBHandler(name)
with context as resource:
    for i in range(5):
        resource.write(str(i))

原则上，... if ... else ... 也可以内联到 with 语句中。为了便于阅读，它被分开了。

使用上下文管理器选择数据库或文件处理程序

Using a context manager to choose db or file handler

python

contextmanager

python-3.x