Python循环中可选写入文件的函数更简洁实现方法

我编写了一个迭代计算量X,Y并返回最终结果的函数，该函数会将每轮迭代得到的X,Y保存到文件中，基础代码结构如下：

def myFunc():
    X,Y = 0,0
    file1 = open(output1,"w")
    file2 = open(output2,"w")
    for i in range(1000):
        X,Y = someCalculation(X,Y) #calculations happen here
        file1.write(X)
        file2.write(Y)
    file1.close()
    file2.close()
    return X,Y

现需求为：调用函数时若省略文件名参数output1或output2，函数仍执行完全一致的计算逻辑，但不对未传入路径的对应文件执行写入操作。

我目前实现的写法较为冗余，代码如下：

def myFunc(output1=None,output2=None):
    X,Y = 0,0
    if (output1 != None): file1 = open(output1,"w")
    if (output2 != None): file2 = open(output2,"w")
    for i in range(1000):
        X,Y = someCalculation(X,Y) #calculations happen here
        if (output1 != None): file1.write(X)
        if (output2 != None): file2.write(Y)
    if (output1 != None): file1.close()
    if (output2 != None): file2.close()
    return X,Y

请问是否存在更优、更整洁的代码编写方式来实现该需求？

你原来的写法除了分支判断重复冗余，还存在资源泄漏风险：如果someCalculation运行中途抛出异常，已经打开的文件对象不会执行close()，文件句柄会被占用直到被GC回收。可以参考下面两种更整洁的实现，完全不需要在循环内写重复判断：

推荐方案：上下文管理器+空写入占位 #

用自定义的可选打开上下文管理器，未传路径时直接返回吞掉所有写入的空IO对象，用with语句自动处理文件关闭，哪怕计算报错也能安全释放资源：

import os
import io
from contextlib import contextmanager

@contextmanager
def optional_open(path, mode="w"):
    if path is None:
        # Python3.7+可直接用io.NullIO()，低版本兼容用系统devnull
        yield io.NullIO()
    else:
        with open(path, mode) as f:
            yield f

def myFunc(output1=None, output2=None):
    X, Y = 0, 0
    with optional_open(output1, "w") as f1, optional_open(output2, "w") as f2:
        for _ in range(1000):
            X, Y = someCalculation(X, Y)
            f1.write(X)
            f2.write(Y)
    return X, Y

轻量方案：统一初始化+try/finally #

如果不想引入额外的工具函数，可以在初始化阶段统一处理所有文件对象，用try/finally保证资源释放，循环内零分支判断：

import os

def myFunc(output1=None, output2=None):
    X, Y = 0, 0
    # 提前完成所有文件对象初始化，未传路径的直接指向系统空设备
    f1 = open(output1, "w") if output1 is not None else open(os.devnull, "w")
    f2 = open(output2, "w") if output2 is not None else open(os.devnull, "w")
    try:
        for _ in range(1000):
            X, Y = someCalculation(X, Y)
            f1.write(X)
            f2.write(Y)
    finally:
        f1.close()
        f2.close()
    return X, Y

小提示：判断None时推荐用is None/is not None而不是==/!=，避免自定义了__eq__方法的对象和None比较时出现误判。

如果后续需要扩展支持更多输出变量，还可以把输出路径存在字典里批量处理，连重复的文件对象定义都能省掉，扩展性更好。

内容的提问来源于stack exchange，提问作者J-J