在Python中实现可选的上下文管理器

Posted 2021-04-19

我们有一个具有以下使用模式的代码库：

factory = DataFactory(args)
dataset = factory.download_and_cache_big_dataset(key)
metadata = dataset.get_some_metadata()

目前，download_and_cache_big_dataset从S3获取一个非常大的文件并将其放在某处。除此之外，确实如此

filename = get_s3_key(key)
filepath = os.path.join(get_tmp_dir(), filename)
s3.download_file(key, filepath)
return BigFileClass(filepath) # gets stored in a class somewhere

但是，此文件不会被删除。当谨慎地调用此函数并依赖于文件缓存时，这很好，但是当它被重复调用并且我们不想填满磁盘时会很糟糕。有没有办法用context manager重构代码，以便我们可以使用它

factory = DataFactory(args)
with factory.download_and_cache_big_dataset(key) as dataset:
    metadata = dataset.get_some_metadata()
    # do something with metadata

# file gets automatically deleted

但关键的是，在不破坏现有用法的情况下，其他代码的工作原理是什么？或者是否需要一个返回上下文管理器的不同方法？

答案

由于您返回BigFileClass的实例来处理/表示数据，我建议如下。

我假设数据文件对每个实例都是唯一的。

• 将实例变量添加到BigFileClass以跟踪数据文件的路径。
• 将__del__方法添加到BigFileClass中，其中删除了数据文件。

编辑：如果您想使用BigFileClass作为上下文管理器，请为__enter__定义__exit__和BigFileClass方法。在这种情况下，__enter__唯一需要做的就是return self。

我会将删除文件的任务留给__del__方法（当BigFileClass的引用计数达到0时）。当您已经删除了数据文件时，让类实例仍然存在感觉不对。

---

备注w.r.t.建筑。

使用工厂对我来说似乎是一种不必要的复杂问题。 IMO，download_and_cache_big_dataset可能只是一个返回BigFileClass实例的函数。