I\'m using python\'s ftplib
to write a small FTP client, but some of the functions in the package don\'t return string output, but print to stdout
.
Here's another take on this. contextlib.redirect_stdout
with io.StringIO()
as documented is great, but it's still a bit verbose for every day use. Here's how to make it a one-liner by subclassing contextlib.redirect_stdout
:
import sys
import io
from contextlib import redirect_stdout
class capture(redirect_stdout):
def __init__(self):
self.f = io.StringIO()
self._new_target = self.f
self._old_targets = [] # verbatim from parent class
def __enter__(self):
self._old_targets.append(getattr(sys, self._stream)) # verbatim from parent class
setattr(sys, self._stream, self._new_target) # verbatim from parent class
return self # instead of self._new_target in the parent class
def __repr__(self):
return self.f.getvalue()
Since __enter__ returns self, you have the context manager object available after the with block exits. Moreover, thanks to the __repr__ method, the string representation of the context manager object is, in fact, stdout. So now you have,
with capture() as message:
print('Hello World!')
print(str(message)=='Hello World!\n') # returns True