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.

Was it helpful?

Solution

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

OTHER TIPS

Use a helper function to wrap the real open() which either calls the real open() or which returns an object that has the methods write(), flush() and close():

class MockFile(object):
    def write(self, data): pass
    def flush(self): pass
    def close(self): pass

def optionalWrite(filename, mode):
    if writeForRead: # <--- Your condition here
        return open(filename, mode)

    return MockFile()

with optionalWrite as output:
    ...

Sounds like you need to pass in a function into the scope which encapsulates the fact that you might or might not want to store the results into a file. In OO languages this would be called a Strategy pattern but in Python you can just pass in a function (since functions are first class.)

my_gen = (compute_stuff(item) for item in items)
results = store_strategy(my_gen)
return results

Wherein store_strategy could just be something that already has the with statement within it or not.

def pass_through(iterable):
    return iterable

def file_store(filename):
    def store(items):
        with open(filename, 'w') as output:
            results = []
            for item in items:
                write_result(item, result, output)
                result.append(item)
        return results
    return store

Most, if not all, of the other answers describe how to write a context manager that would allow you to do what you want. Here's something a little different that directly addresses your question:

Can a “with” statement be used conditionally?

Yes, it can — by using a generator function in conjunction with a for loop that conditionally doesn't iterate. Here's a runnable example based on the code in your question showing what I mean:

# Scaffolding added for testing.
def compute_stuff(item):
    return 'item: ' + str(item)

def write_result(item, result, output):
    output.write(result + '\n')

# Generator function.
def conditionally_with(filename, mode='r'):
    if not filename:  # Check condition.
        return
    else:
        with open(filename, mode) as opened:
            yield opened

if __name__ == '__main__':

    filename = 'foobar.txt'
    items = range(5)
    results = []

    for output in conditionally_with(filename, "w"):
        for item in items:
            result = compute_stuff(item)
            write_result(item, result, output)
            results.append(result)

    print(results)  # -> ['item: 0', 'item: 1', 'item: 2', 'item: 3', 'item: 4']
    # return results

Using coroutines

http://www.dabeaz.com/coroutines/Coroutines.pdf (as suggested by Paulo Scardine)

If we want to write:

def writer(filename):
  with open(filename, "w") as output:
    while True:
      try:
        item, result = (yield)
        write_result(item, result, output)
      except GeneratorExit:
        output.flush()
        break

If we don't:

def dummy_writer():
   while True:
     yield

Initialize our coroutine:

result_writer = writer(filename) if filename else dummy_writer()
result_writer.next()

Run our code:

results = []
for item in items:
    result = compute_stuff(item)
    result_writer.send((item, result))
    results.append(result)
result_writer.close()
return results

Here's something derived from the suggestion in wheaties' answer which I think may be the best context-manager-free approach (and so deserves example code which illustrates it more concretely):

def create_list():
    return list

def file_store(filename, mode='w'):
    def store(items):
        with open(filename, mode) as output:
            results = []
            for item in items:
                write_result(item, output)
                results.append(item)
        return results
    return store

store_strategy = file_store(filename) if filename else create_list()
results = store_strategy(compute_stuff(item) for item in items)
return results

What you are trying to do is cause late creation of the file. What you want is something that looks like a context manager, but doesn't actually create the file until you need it. You'll need to implement __enter__ and __exit__ yourself. This is a (very) abbreviated example, compare to a full one, just for your exact case:

class LateFile(object):
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.fp = None

    def __enter__(self):
        # Do nothing - we don't open the file yet
        return(self)

    def __exit__(self, exctype, value, traceback):
        if self.fp != None:
            fp.close()

    def write(self, *args, **kwargs):
        if self.fp == None:
            self.fp = open(self.filename, self.mode)
        self.fp.write(*args, **kwargs)

Something along those lines.

And then, to use it, do something like:

with LateFile(filename, "w") as output:
    for item in items:
        result = compute_stuff(item)
        if should_write_result(item, result):
            write_result(item, result, output)
        results.append(result)

write_result should see output as a normal file object; you'll need to reflect or delegate the methods down to it. Doing this, if no results are written, the file will not be created, but if even one result is written, the file will be created & written to.

Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top