In a loop, optionally write output to files-CodePudding

I wrote a function which iteratively computes some quantities X,Y, returning the final result. Additionally, this code saves each iteration of X,Y to a file. Here is the basic structure:

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

However, if the filename output1 or output2 is omitted when the function is called, I need this function to perform the same calculation without appending anything to the relevant file.

Here is my messy solution:

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

Is there a better, cleaner way to write this?

CodePudding user response：

Make a dummy file object that ignores writes, and supports the context manager interface:

class NoFile:
    def __enter__(self): return self
    # Propagate any exceptions that were raised, explicitly.
    def __exit__(self, exc_type, exc_value, exc_tb): return False
    # Ignore the .write method when it is called.
    def write(self, data): pass
    # We can extend this with other dummy behaviours, e.g.
    # returning an empty string if there is an attempt to read.

Make a helper function that creates one of these instead of a normal file when the filename is None:

def my_open(filename, *args, **kwargs):
    return NoFile() if filename is None else open(filename, *args, **kwargs)

Use with blocks to manage the file lifetimes, as you should do anyway - but now use my_open instead of open:

def myFunc(output1=None,output2=None):
    X, Y = 0, 0
    with my_open(output1, 'w') as f1, my_open(output2, 'w') as f2:
        for i in range(1000):
            X, Y = someCalculation(X, Y) #calculations happen here
            f1.write(X)
            f2.write(Y)
    return X, Y

CodePudding user response：

You can create a dummy class that has a do-nothing write method. ExitStack is used to ensure any opened files are closed automatically.

from contextlib import ExitStack


class NoWrite:
    def write(self, value):
        pass


def myFunc(output1=None, output2=None):
    X,Y = 0,0
    with ExitStack() as es:
        file1 = es.enter_context(open(output1, "w")) if output1 is not None else NoWrite()
        file2 = es.enter_context(open(output2, "w")) if output2 is not None else NoWrite()

        for i in range(1000):
            X,Y = someCalculation(X, Y) #calculations happen here
            file1.write(X)
            file2.write(Y)
    return X,Y

As you appear to be logging the X and/or Y values at each step, you may want to look into using the logging module instead, creating a FileHandler for the appropriate logger instead of passing output file names to myFunc itself.

import logging

# configuration of the logger objects is the responsibility
# of the *user* of the function, not the function itself.
def myFunc():
    x_logger = logging.getLogger("myFunc.x_logger")
    y_logger = logging.getLogger("myFunc.y_logger")

    X,Y = 0,0
    for i in range(1000):
        X,Y = someCalculation(X, Y)
        x_logger.info("%s", X)
        y_logger.info("%s", Y)
    return X,Y