Can I redirect the stdout in python into some sort of string buffer?

前端 未结 9 1560
野趣味
野趣味 2020-11-22 06:23

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.

9条回答
  •  悲哀的现实
    2020-11-22 06:54

    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
    

提交回复
热议问题