python 将函数标准输出存至变量问题

背景

在编写一个ansible模块，调用一个开源库的时候，遇到一个问题：这个库中的函数都是将结果输出到“标准输出”或“标准错误”的，那么如何将这些内容赋值给变量，方便操作呢？

那是否可以修改下这个库，让其支持return呢？当然是可以的，不过，好麻烦啊，因为好多函数，你用哪个，就得修改哪个。这不可取。

能不用这个库，选用别的库吗？当然可以了。但是这里讨论的问题，已经不是解决最初的编写模块的问题了。

现在的问题就是：如何将函数或方法的标准输出内容存至变量当中？

系统环境

• OS: CentOS-7.9
• python: 3.8.5

使用重定向

这时候就想到了重定向，输出是一种io流，貌似只能重定向。

重定向至文件

import io
import sys

# 假设这个函数没办法修改，只能调用
def stdout_func(name):
    print("hello", name)

def redirect_to_file(func, *args, **kw):
    file_path = '/tmp/.test_io'
    # 这种低级别的文件对象，只能写入bytes
    f_w_io = io.FileIO(file_path, 'w')
    f_r_io = io.FileIO(file_path, 'r')
    # 将其包装成可读写 str 的
    f_w_io = io.TextIOWrapper(f_w_io)
    f_r_io = io.TextIOWrapper(f_r_io)
    # 保存旧的 标准输出流
    old_stdout = sys.stdout
    # 将输出流指向创建好的文件对象
    sys.stdout = f_w_io
    # 调用函数
    func(*args, **kw)
    # 刷新至文件
    f_w_io.flush()
    # 关闭可写对象
    f_w_io.close()
    # 读取数据
    value = f_r_io.read()
    # 关闭可读对象
    f_r_io.close()
    # 将标准输出置为原先值
    sys.stdout = old_stdout
    return value
    
if __name__ == '__main__':
    v1 = redirect_to_file(stdout_func, "swk")
    print("v1", v1)

脚本运行，可以获得v1 hello swk, print函数会自动添加一个换行符，这个自己处理下即可。

或者直接使用open函数打开的file-like Object。效果一样，

使用StringIO

def redirect_to_stringio(func, *args, **kw):
    output = io.StringIO()
    # 保存旧的 标准输出流
    old_stdout = sys.stdout
    # 将输出流指向创建好的StringIO对象
    sys.stdout = output
    # 调用函数
    func(*args, **kw)
    # 读取数据
    value = output.getvalue()
    # 关闭可读对象
    output.close()
    # 将标准输出置为原先值
    sys.stdout = old_stdout
    return value

效果与文件一样，但是这个是在内存中操作，相对来说，效率更高。

实战中翻车

有了上面的之后，开始实战。很快就翻车了。

要调用的开源库的函数，长这个样子。

def format_output(value):
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    stream = sys.stdout
    print("format_ouatput", value)
    stream.flush()

报错'_io.StringIO' object has no attribute 'buffer'

由于format_output函数中有sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')。

而StringIO对象的确是没有buffer属性的，所以就报错了。

在python官方文档有如下说明

Note: To write or read binary data from/to the standard streams, use the underlying binary buffer object. For example, to write bytes to stdout, use sys.stdout.buffer.write(b’abc’).

However, if you are writing a library (and do not control in which context its code will be executed), be aware that the standard streams may be replaced with file-like objects like io.StringIO which do not support the buffer attribute.

sys.stdout.buffer对象

sys.stdout.buffer is the underlying binary buffer object.

字面意思：底级别二进制的缓冲对象。

实际效果，就是其不可以直接写入字符串。

而使用了io.TextIOWrapper对其包装之后，其就可以写入字符串了，就好像在redirect_to_file函数中使用的那样。

同时也指定了其字符集。

而且其还拥有了buffer效果, 即在写入数据后，不会马上写入，而是先存储到buffer当中，需要flush()之后才能看到效果。演示如下:

In [36]: print("hello world")
hello world

In [37]: sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
In [38]: print("hello buffer")
In [39]: print("where are you?")
In [40]: sys.stdout.flush()

hello buffer

where are you?

解决

为StringIO对象添加一个buffer属性

由于format_output这个被调用函数，总是要使用sys.stdout.buffer属性，所以这里就需要为StringIO添加这样一个属性。

上面看到了sys.stdout.buffer is the underlying binary buffer object。

io.BytesIO是一个二进制对象流，尝试一下。

def redirect_to_stringio(func, *args, **kw):
    output = io.StringIO()
    # 生成一个 BytesIO对象
    b_io = io.BytesIO()
    # 保存旧的 标准输出流
    old_stdout = sys.stdout
    # 将输出流指向创建好的StringIO对象
    output.buffer = b_io
    sys.stdout = output
    # 调用函数
    func(*args, **kw)
    # 读取数据
    b_value = b_io.getvalue()
    b_value = b_value.decode('utf-8')
    # 关闭可读对象
    output.close()
    b_io.close()
    # 将标准输出置为原先值
    sys.stdout = old_stdout
    return b_value

这样改好之后，输出就正常了。

使用io.StringIO当作buffer属性

这种情况下，会报错TypeError: string argument expected, got 'bytes'。

使用io.BufferedWriter(StringIO)对象当作buffer属性

这种情况下，会报错TypeError: string argument expected, got 'memoryview'。

之所以尝试这个，是因为sys.stdout.buffer是一个_io.BufferedWriter对象: <_io.BufferedWriter name='<stdout>'>

使用io.FileIO当作buffer属性

b_io = io.FileIO('/tmp/.test_io', 'wb')

这种也是可以的，不过同时还得再打开一个file对象，读取数据。

def redirect_to_stringio(func, *args, **kw):
    file_path = '/tmp/.test_io'
    output = io.StringIO()
    # 生成一个 FileIO对象
    b_io = io.FileIO(file_path, 'wb')
    b_io_r = io.FileIO(file_path, 'rb')
    # 保存旧的 标准输出流
    old_stdout = sys.stdout
    # 将输出流指向创建好的StringIO对象
    output.buffer = b_io
    sys.stdout = output
    # 调用函数
    func(*args, **kw)
    # 显式flush()，保证数据落盘
    b_io.flush()
    # 读取数据
    b_value = b_io_r.read()
    b_value = b_value.decode('utf-8')
    # 关闭可读对象
    b_io.close()
    b_io_r.close()
    output.close()
    # 将标准输出置为原先值
    sys.stdout = old_stdout
    return b_value