Como lidar bem com `with open (…)` e `sys.stdout`?

Freqüentemente, preciso enviar dados para o arquivo ou, se o arquivo não for especificado, para o stdout. Eu uso o seguinte snippet:

if target:
    with open(target, 'w') as h:
        h.write(content)
else:
    sys.stdout.write(content)

Gostaria de reescrevê-lo e lidar com os dois alvos de maneira uniforme.

No caso ideal, seria:

with open(target, 'w') as h:
    h.write(content)

mas isso não funcionará bem porque sys.stdout é fechado ao sair do withbloco e eu não quero isso. Eu nem quero

stdout = open(target, 'w')
...

porque eu precisaria me lembrar de restaurar o stdout original.

Relacionado:

• Redirecionar stdout para um arquivo em Python?
• Manipulando exceções - artigo interessante sobre como lidar com exceções em Python, em comparação com C ++

Editar

Eu sei que posso encapsular target, definir funções separadas ou usar o gerenciador de contexto . Procuro uma solução simples, elegante e idiomática que não exija mais do que 5 linhas

Apenas pensando fora da caixa aqui, que tal um open()método personalizado ?

import sys
import contextlib

@contextlib.contextmanager
def smart_open(filename=None):
    if filename and filename != '-':
        fh = open(filename, 'w')
    else:
        fh = sys.stdout

    try:
        yield fh
    finally:
        if fh is not sys.stdout:
            fh.close()

Use-o assim:

# For Python 2 you need this line
from __future__ import print_function

# writes to some_file
with smart_open('some_file') as fh:
    print('some output', file=fh)

# writes to stdout
with smart_open() as fh:
    print('some output', file=fh)

# writes to stdout
with smart_open('-') as fh:
    print('some output', file=fh)

Fique com seu código atual. É simples e você pode dizer exatamente o que está fazendo apenas olhando para ele.

Outra maneira seria com um embutido if:

handle = open(target, 'w') if target else sys.stdout
handle.write(content)

if handle is not sys.stdout:
    handle.close()

Mas isso não é muito mais curto do que o que você tem e parece provavelmente pior.

Você também pode tornar sys.stdoutimpossível fechar, mas isso não parece muito pitônico:

sys.stdout.close = lambda: None

with (open(target, 'w') if target else sys.stdout) as handle:
    handle.write(content)

Por que LBYL quando você pode EAFP?

try:
    with open(target, 'w') as h:
        h.write(content)
except TypeError:
    sys.stdout.write(content)

Por que reescrevê-lo para usar o bloco with/ asuniformemente quando você precisa fazê-lo funcionar de uma maneira complicada? Você adicionará mais linhas e reduzirá o desempenho.

Outra solução possível: não tente evitar o método de saída do gerenciador de contexto, apenas duplique o stdout.

with (os.fdopen(os.dup(sys.stdout.fileno()), 'w')
      if target == '-'
      else open(target, 'w')) as f:
      f.write("Foo")

Uma melhoria da resposta de Wolph

import sys
import contextlib

@contextlib.contextmanager
def smart_open(filename: str, mode: str = 'r', *args, **kwargs):
    '''Open files and i/o streams transparently.'''
    if filename == '-':
        if 'r' in mode:
            stream = sys.stdin
        else:
            stream = sys.stdout
        if 'b' in mode:
            fh = stream.buffer  # type: IO
        else:
            fh = stream
        close = False
    else:
        fh = open(filename, mode, *args, **kwargs)
        close = True

    try:
        yield fh
    finally:
        if close:
            try:
                fh.close()
            except AttributeError:
                pass

Isso permite E / S binária e passa eventuais argumentos estranhos para opense filenamerealmente for um nome de arquivo.

Eu também escolheria uma função de wrapper simples, que pode ser muito simples se você puder ignorar o modo (e, consequentemente, stdin x stdout), por exemplo:

from contextlib import contextmanager
import sys

@contextmanager
def open_or_stdout(filename):
    if filename != '-':
        with open(filename, 'w') as f:
            yield f
    else:
        yield sys.stdout

Ok, se estamos entrando em guerras de uma linha, aqui está:

(target and open(target, 'w') or sys.stdout).write(content)

Gosto do exemplo original de Jacob, contanto que o contexto seja escrito apenas em um lugar. Seria um problema se você acabasse reabrindo o arquivo para muitas gravações. Acho que só tomaria a decisão uma vez no início do script e deixaria o sistema fechar o arquivo ao sair:

output = target and open(target, 'w') or sys.stdout
...
output.write('thing one\n')
...
output.write('thing two\n')

Você pode incluir seu próprio manipulador de saída se achar que é mais organizado

import atexit

def cleanup_output():
    global output
    if output is not sys.stdout:
        output.close()

atexit(cleanup_output)

Se você realmente deve insistir em algo mais "elegante", ou seja, uma linha:

>>> import sys
>>> target = "foo.txt"
>>> content = "foo"
>>> (lambda target, content: (lambda target, content: filter(lambda h: not h.write(content), (target,))[0].close())(open(target, 'w'), content) if target else sys.stdout.write(content))(target, content)

foo.txtaparece e contém o texto foo.

Que tal abrir um novo fd para sys.stdout? Desta forma, você não terá problemas para fechá-lo:

if not target:
    target = "/dev/stdout"
with open(target, 'w') as f:
    f.write(content)

if (out != sys.stdout):
    with open(out, 'wb') as f:
        f.write(data)
else:
    out.write(data)

Melhora leve em alguns casos.

import contextlib
import sys

with contextlib.ExitStack() as stack:
    h = stack.enter_context(open(target, 'w')) if target else sys.stdout
    h.write(content)

Apenas duas linhas extras se você estiver usando Python 3.3 ou superior: uma linha para o extra importe uma linha para o stack.enter_context.