Python如何获取subprocess子进程返回的变量并在主程序中使用

问题背景 #

运行Python程序main.py时，需要调用子脚本sub.py，同时满足核心要求：一是sub.py运行时的ANSI转义样式（如粗体、红字）能正常在终端渲染，二是main.py可以直接获取sub.py运行后返回的变量z。原先使用subprocess调用的方案存在无法获取返回值、或者ANSI样式丢失的问题。

原始代码 #

sub.py #

import optparse
#ANSI escape sequence
BOLD = "\033[1m"
RED = "\033[1;31m"
END = "\033[0m"

def parser():
    parser = optparse.OptionParser()
    parser.add_option(
        "-n", help="Enter a number: ", dest='n')
    (options, args) = parser.parse_args()
    n = options.n
    return n

def main():
    n = parser()
    print(BOLD + RED + f'Your number is {n}' + END)
    z = int(n)**2
    return z

if __name__ == "__main__":
    z = main()

main.py #

import subprocess
# 该方法仅在终端运行带ANSI转义序列的sub.py，无法返回变量z
subprocess.run('python /home/username/sub.py -n 2', shell=True)

# 该方法运行sub.py时终端输出无ANSI转义效果，同样无法返回变量z
output = subprocess.check_output('python /home/username/sub.py -n 2', shell=True) 
print(output)

最优实现方案（无需subprocess） #

直接将sub.py作为模块导入main.py调用是性能最高、实现最简单的方案，完全满足所有需求：ANSI转义序列会直接输出到当前终端保留样式，同时可以直接拿到函数返回值。你自己给出的修改方案就是正确的实现，优化后代码如下：

修改后的sub.py #

import optparse
#ANSI escape sequence
BOLD = "\033[1m"
RED = "\033[1;31m"
END = "\033[0m"

def parser():
    parser = optparse.OptionParser()
    parser.add_option(
        "-n", help="Enter a number: ", dest='n')
    (options, args) = parser.parse_args()
    n = options.n
    return n

def for_main(n):
    print(BOLD + RED + f'Your number is {n}' + END)
    z = int(n)**2
    return z

def main():
    n = parser()
    print(BOLD + RED + f'Your number is {n}' + END)
    z = int(n)**2
    return z

if __name__ == "__main__":
    main()

修改后的main.py #

from sub import parser
from sub import for_main
import optparse

#ANSI escape sequence
BOLD = "\033[1m"
RED = "\033[1;31m"
END = "\033[0m"

def main():
   n = parser()
   z = for_main(n)

if __name__ == "__main__":
    main()

特殊场景方案（必须使用subprocess调用） #

如果因为权限、隔离要求等原因必须用subprocess启动sub.py，可以用以下方法同时保留ANSI效果和获取返回值：

1. 在sub.py的main函数最后新增一行单独打印返回值z
2. main.py中调用subprocess时同时捕获输出、并将输出实时打印到终端，最后提取最后一行的z值

示例代码如下：

调整后的sub.py新增逻辑 #

if __name__ == "__main__":
    z = main()
    print(f"RETURN_VALUE:{z}") # 新增单独打印返回值的行

调整后的main.py #

import subprocess

proc = subprocess.Popen(
    'python /home/username/sub.py -n 2',
    shell=True,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True
)

output_lines = []
for line in proc.stdout:
    print(line, end='') # 实时打印输出，保留ANSI效果
    output_lines.append(line.strip())

# 提取返回值
z_line = [l for l in output_lines if l.startswith("RETURN_VALUE:")][0]
z = int(z_line.split(":")[1])

内容的提问来源于stack exchange，提问作者Neal Titus Thomas