如何定義可以與with語句一起使用的Python類？

Question

我知道StringIO就像一個文件對象，可以open('somefile.txt')你從open('somefile.txt') 。

現在我想使用帶有with語句的StringIO ：

with StringIO('some string') as fh: # fh as in "file handle"
    data = [stuff from stuff in fh.read()]

但Python抱怨StringIO類型沒有__exit__方法。 StringIO ：

class MyStringIO(StringIO):
    def __exit__(self):
        self.close()

我現在得到一個關於沒有__enter__方法的例外。 如何定義__enter__方法？ Python對可以與with語句一起使用的類有什么期望？

Answer 1

您需要編寫上下文管理器 。 如果您不想編寫整個協議，可以使用contextlib.contextmanager裝飾器簡化它。

Answer 2

只是舉一個例子來說它是矯枉過正的，請使用contextlib，它更適用於生成器：

這是一個很好的有用的上下文管理器示例。 我在David Beazley的博客中找到了它：

http://dabeaz.blogspot.de/2010/02/context-manager-for-timing-benchmarks.html

並且只是稍微修改它以使用每個平台的最佳計時器，這是我在timeit模塊中使用的技巧：

# benchmark.py
import timeit
class benchmark(object):
    from timeit import default_timer as timer
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        self.start = self.timer()
    def __exit__(self, ty, val, tb):
        end = self.timer()
        print("%s : %0.3f seconds" % (self.name, end-self.start))
        return False

現在，它為您提供了一個上下文管理器，您可以像這樣使用：

with benchmark('algorithm 1'):
    ... (some computation)
with benchmark('other approach'):
    ... (some other computation)

並打印出你的輸出。 我喜歡這個例子，想分享它。

Answer 3

您可以像使用__exit__一樣定義__enter__方法。 這些是對象與with語句一起使用所需的兩種方法，有關詳細信息，請參閱此方法。