巧妙使用 Python StringIO 拦截控制台输出

开头说两句

• StringIO 实际上就是在内存中读写 str 类型的数据。
• 在这个 demo 中还会用到关于 sys 库的输出 stdout 模块。(拦截 stdin 和 stderr 也是同理)
  • stdin 用于所有交互式输入（包括调用 input()）;
  • stdout 用于输出 print() 和表达式语句以及提示 input();
  • stderr 用于将解释器自己的提示及其错误消息输出。

进入 demo

• demo 代码

  import sys
  from io import StringIO
  
  
  class ConsoleOutputRedirect:
      """ Wrapper to redirect stdout or stderr """
      def __init__(self, fp):
          self.fp = fp
  
      def write(self, s):
          self.fp.write(s)
  
      def writelines(self, lines):
          self.fp.writelines(lines)
  
      def flush(self):
          self.fp.flush()
  
  
  stdout_redirect = ConsoleOutputRedirect(sys.stdout)
  
  
  def print_something():
      # should be print something but it shouldn`t
      print("Hello Leo!")
  
  
  def block_output():
      # Block Console stdout
      stdout_redirect.fp = StringIO()
      temp_stdout, sys.stdout = sys.stdout, stdout_redirect
  
      # Core Method
      print_something()
  
      # Swap sys.stdout
      sys.stdout = temp_stdout
  
      # Block output
      print("Today is Good Weather, {}".format(stdout_redirect.fp.getvalue()))
  
  
  if __name__ == '__main__':
      block_output()
  

• 运行完这段代码后会惊奇的发现控制台中只输出了 Today is Good Weather, Hello Leo!。（这就是拦截的魔力）

demo 讲解

• 大体的讲解:
  • 实际上就是将需要拦截部分的 sys.stdout (即需要拦截 print 的内容) 通过自定义 StringIO 的方式进行截获，最终再通过其他方式将拦截内容呈现出来即可。
• 讲解点：
  • 1、首先得了解 Python 中 print 这个函数在输出的过程中干了哪些内容。
    • 官方文档 中也提到，实际上 print 在输出的过程使用的就是输入输出流。print 中的内容通过 write stream 的方式写入到 str 中。而正好在 Python 中管理 IO 的类就是 _IOBase, 也是 StringIO 的父类。
  • 2、拦截输出的核心是基于 StringIO 的这个类读取输出的内容，因此就需要了解 StringIO 关于输出部分的核心方法。StringIO, 以及拦截所需要的方法。
    • 拦截过程前首先得让 print 或者叫 sys.stdout 的内容赋予给我们代码中的一个变量，这样我们才能对这个输出内容进行拦截或者叫 hack。
    • 拦截之后，将空的输出流赋给 sys.stdout 即可，代码中是在拦截之前将一个未使用的 sys.stdout 输出流赋给了一个临时变量，最后再通过 StringIO 的 getvalue 从我们自定义的 StringIO 对象中取出我们需要的内容就完成了。
  • 3、实际上就是拦截前和拦截后两部分的 sys.stdout 的交换。（大体上可以理解成两个变量交换[略微不严谨的理解]）