CODEWITHSUNDEEP
pythoncontextmanager
Ethan Posner
Ethan Posner

Reputation: 386

Use Conditional within Context Manger

I want to be able to do something like this:

file1 = 'data/f1.txt'
file2 = None
file3 = 'data/f3.txt'
with open(file1) as f1, open(file2) as f2, open(file3) as f3:
    for i in range(1000):
        x,y,z = func(i)
        if f1: f1.write(f"{x}\n")
        if f2: f2.write(f"{y}\n")
        if f3: f3.write(f"{z}\n")

So that output is only written to files if paths are provided.

I tried the following:

file1 = 'data/f1.txt'
file2 = None
file3 = 'data/f3.txt'
with open(file1) if file1 else None as f1, open(file2) if file2 else None as f2, open(file3) if file3 else None as f3:
    for i in range(1000):
        x,y,z = func(i)
        if f1: f1.write(f"{x}\n")
        if f2: f2.write(f"{y}\n")
        if f3: f3.write(f"{z}\n")

But I got TypeError: expected str, bytes or os.PathLike object, not NoneType.

Is there a simple, pythonic way to achieve this?

Thanks.

Upvotes: 0

Views: 66

Answers (1)

J_H
J_H

Reputation: 20450

By the time you've invented identifiers file{1,2,3}, that's starting to be a Code Smell that you really want to fill a container with files.

Create a class FilesWriter that is a context manager. Pass in a bunch of filespecs. It will simply discard any empty ones:

    self.fspec_to_file = {}
    for fspec in fspecs:
        if fspec is not None:
            self.fspec_to_file[fspec] = open(fspec, 'w')

Now your instance is responsible for closing those when the with calls your __exit__ method.

In the body of the with, caller will either invoke a write_to_all(s: str) method you provide, or will use the elements of fspec_to_file as needed.

So calling sequence would look like:

    with FilesWriter([fspec1, fspec2, fspec3]) as fw:
        for i in range(1000):
            result = compute()
            fw.write_to_all(result)

and upon leaving the block your __exit__ method does:

    for f in self.fspec_to_file.values():
        f.close()

Strictly speaking .close() can fail, e.g. due to full disk. Put it in a try block if needed.

If any open() fails, it might make sense to close all previously opened files and die with fatal error.

You seemed to need to know details about each open file. But if not, feel free to convert that dict to a simple set of open file handles.

Upvotes: 2

Related Questions