How would you rewrite this Python code?

I ran across this problem while writing some code for a challenge.

I have some code that I want to write to a particular file if a filename is provided, or standard output if not:

with (open(filename, 'w') if filename else sys.stdout) as file:
    do_something(file)

That with line is a bit long, and doesn't seem readable at a glance. I could put the open outside the with:

f = open(filename, 'w') if filename else sys.stdout
with f as file:
    do_something(file)

The file should be closed when the with is exited. I'm done with stdout in this case, but I could just have well have wanted to use it for other things later on in the program. But the worst thing about this way of doing it seems to be that making a habit of using open outside a with expression could lead to forgetting to close files.

I could go the try/finally route:

try:
    f = open(filename, 'w') if filename else sys.stdout
    do_something(f)
finally:
    if f is not sys.stdout:
        f.close()

But this seems a bit verbose and reminds me too much of Java, which, as a rule of thumb, probably means it's not Pythonic.

I could write a context manager to hide the "check if it's stdout and close it" logic:

from contextlib import contextmanager

@contextmanager
def file_or_stdout(fname):
    if fname:
        f = open(fname, 'w')
        yield f
        f.close()
    else:
        yield sys.stdout

with file_or_stdout(filename) as file:
    do_something(file)

This seems pretty clear, but it also might be overkill. I'm leaning toward this though, as it leaves the do_something block plenty of room for clarity, and I could extract the file_or_stdout function into a utility library so it's not cluttering up the file.

Any thoughts?

Comments

[rhymes]

I think the simplest is the first one :D

Or you can wrap that if in a normal function if you plan to use it more than once:

def file_or_stdout(fname=None):
  if fname:
    return open(fname, 'w')
  else:
    return sys.stdout

I don't think you need to use a context manager in this case because there are no other operations between the opening and closing of the file.

Anyway I wouldn't do it like that.

Keep in mind that if you use sys.stdout this way you will close it which means that any attempt to write on it after your with statement will result in an error:

>>> import sys
>>> with file_or_stdout() as f:
...     f.write("Yo!\n")
...
Yo!
4
>>> sys.stdout.write("Yo 2\n")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: I/O operation on closed file.
>>> print("yo 3")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: I/O operation on closed file.

So why do you need to hide stdout this way? It's probably easier to use StringIO, write to it and then use a print to display it on stdout.

[Dustin King]

So why do you need to hide stdout this way?

I'm just trying to keep the code clear and non-repetitive. In terms of level of effort and maintainability, if it needs to be able to optionally write to a file, breaking the inner logic out into a function as inspired by @thinkslynk is probably best.

But if it's a one-off project it might be even better to only write to stdout and pipe it to a file in the shell if needed. (Sometimes it's fun to over-engineer though.)

[Tomer]

Actually the context manager solves this by only opening/closing the fp if it isn't to stdout so it might be the best solution

[rhymes]

Oh yes, you're right :-)

+1 to the context manager, though it might be overkill if the function is used only once.

[Stephen Dycus]

I think the cleanest way is this:

if filename:
    with (open(filename, 'w') as file:
        do_something(file)
else:
    do_something(sys.stdout)

[Dustin King]

Certainly works if do_something is just a function call. I meant it as a stand-in for a larger block of code. But maybe that block should actually be a function.

[Niels Böhm]

Yes, it should be a function. stdout/stderr/stdin are different from other streams in that they're managed outside (by the OS and Python, not by your app) whereas an open should be used with with in simple cases if possible.

You can start treating them the same when you start writing to / reading from them and stop when you stop doing that. All that can nicely be stuck away in a function. That also enables you to call the actual I/O part with a different file-like thing, say a StringIO instance when testing.

[Dustin King]

I don't think I'd heard of CQS before. I'll have to think more about it. It does make the code clear in this case.

On first examination, one issue I have with the general concept is that it seems that functions that act like constructors, such as file_or_stdout() in my example, or even just output = open(...) would violate it because they both do something (e.g. open the file, build a context manager) and return a reference to the thing they've opened or created.

CQS seems similar to a rule of thumb that parameters are for input and return values are for output, which I'm violating with do_something(), because it was a stand-in for a couple lines in the original code.

Maybe the larger takeaway of this exercise is that mutability has its pitfalls.

[Derek Bartron]

How is the first one harder to read then the rest that more lines and more complicated? If a Python coder can't read an inline expression, how is turning it into a context manager helping? I'm confused by this entire article.

[Dustin King]

There is a lot of subjectivity involved in deciding which code is "better", and that line may be perfectly readable for you.

To me it seemed like there was too much happening in that one line, which would mean someone has to stop on it for a few seconds to figure out what it's doing. Having a more complicated with clause also seemed to distract from the code after with(...):. This isn't shown very well here, since I just used do_something(file) as a stand-in, instead of the couple lines in the linked code this was derived from.

But you may have a point in that one or two long lines might be preferable to a lot of short lines. The point of this post was that I ran across a situation in which the clearest way to write the code wasn't obvious, and I thought it would be valuable to hear peoples opinions on different ways to do it, and the tradeoffs between them.

[Yoandy Rodriguez Martinez]

I'll tell you as my old Python sensei used to tell me: When in doubt, import this.
I think your second solution is the one. It is both concise and does not messes with PEP 8.

[Dustin King]

What do you mean by "import this"?

[Pierre Bouillon]

Just type "import this" I'm your interpreter, you will understand :)

[Dustin King]

Ah yes, the good ol' Zen of Python.

[Tomer]

TIL building a context manager is so simple and clean.