Make TeeStream context aware.
@@ -723,10 +723,21 @@ class TeeStream(object):
self.consumers.append(consumer)
def __del__(self):
- self.flush()
+ self.close()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self):
+ self.close()
- # close open files
+ def close(self):
+ """
+ Flushes all registered consumer streams and closes opened files.
+ """
+ self.flush()
for f in self.open_files:
+ if not getattr(f, "closed", False):
f.close()
def write(self, *args, **kwargs):
