Source code for pyquickhelper.pycode.utils_tests_stringio

"""
A :epkg:`*py:io:StringIO` which also outputs its content for a file.


:githublink:`%|py|5`
"""
import os
from io import StringIO


[docs]class StringIOAndFile(StringIO): """ Overloads a StringIO to also writes in a file. :githublink:`%|py|12` """
[docs] def __init__(self, filename): """ constructor :param filename: filename :githublink:`%|py|19` """ StringIO.__init__(self) self.filename = filename if os.path.exists(filename): os.remove(filename) self.handle = None self.redirect = {} self.to = None
[docs] def write(self, s): """ Appends a string. :param s: add a string to the stream :return: self :githublink:`%|py|34` """ if self.to is not None: self.redirect[self.to].write(s) else: StringIO.write(self, s) if not self.handle: self.handle = open(self.filename, "w", encoding="utf-8", errors="ignore") self.handle.write(s) self.handle.flush() return self
[docs] def flush(self): """ Calls two flush. :githublink:`%|py|49` """ StringIO.flush(self) if self.handle: self.handle.flush()
[docs] def close(self): """ Calls two close. :githublink:`%|py|57` """ StringIO.close(self) if self.handle: self.handle.close()
[docs] def begin_test(self, name): """ Redirects output to a :epkg:`*py:io:StringIO`. :param name: name :githublink:`%|py|67` """ if self.to is not None: raise Exception("A test has not finished: '{0}'".format(self.to)) if name is None: raise ValueError("name is None") self.to = name self.redirect[name] = StringIO()
[docs] def end_test(self, name): """ Ends the redirection. :param name: name :githublink:`%|py|80` """ if name != self.to: raise ValueError( "Inconsistency in test name '{0}' != '{1}'".format(name, self.to)) self.to = None
[docs] def getvalue(self): """ Returns the content of the buffer. :githublink:`%|py|89` """ if self.to is not None: return self.redirect[self.to].getvalue() else: return StringIO.getvalue(self)