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

Question

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

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语句不需要作为上下文管理器类型的可调用对象——任何计算为上下文管理器的表达式都可以。 这允许 function 调用、条件表达式以及直接名称引用来提供上下文管理器。

您几乎可以删除所有间接，并直接 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语句中。 为了便于阅读，它是分开的。