cai/benchmarks/Prompt2PwnBench/utils/helpers.py

19 lines
398 B
Python
Executable File

class TeeStream:
"""
Class defined to save and display terminal outputs.
"""
def __init__(self, stream1, stream2):
self.stream1 = stream1
self.stream2 = stream2
def write(self, data):
self.stream1.write(data)
self.stream2.write(data)
self.flush()
def flush(self):
self.stream1.flush()
self.stream2.flush()