Can a "with" statement be used conditionally?

Question

I have a Python function of roughly the following structure, that computes some results and writes them to a file:

results = []
with open(filename, "w") as output:
    for item in items:
        result = compute_stuff(item)
        write_result(item, result, output)
        results.append(result)
return results

Now I don't always want to write the results to file - sometimes I only want to compute them and have them returned. Making "write_result" conditional is easy, but is there also a way to make the file creation in the "with" statement condition-dependent? (I know I could handle opening and closing of the file explicitly, but I would incur the "try/finally" overhead that the "with" statement was created to avoid.)

Is there an elegant solution for this?

Edited to add: I may have oversimplified the example. Instead of writing to an arbitrary file, I'm using matplotlib.backends.backend_pdf.PdfPages (link), and adding a plot (new page in the PDF) in every step. In particular, that means I cannot re-open the PDF file with PdfPages, because it would be overwritten.

Answer 1

You can write your own context manager function:

class Dummysink(object):
    def write(self, data):
        pass # ignore the data
    def __enter__(self): return self
    def __exit__(*x): pass

def datasink(filename):
    if filename:
        return open(filename, "w")
    else:
        return Dummysink()

...

results = []
with datasink(filename) as output:
    for item in items:
        result = compute_stuff(item)
        write_result(item, result, output)
        results.append(result)
return results