Come gestire bene sia `con open (...)` che `sys.stdout`?

Spesso ho bisogno di inviare i dati su file o, se il file non è specificato, su stdout. Uso il seguente frammento:

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

Vorrei riscriverlo e gestire entrambi gli obiettivi in ​​modo uniforme.

Nel caso ideale sarebbe:

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

ma questo non funzionerà bene perché sys.stdout è chiuso quando si lascia il withblocco e non lo voglio. Nemmeno io lo voglio

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

perché avrei bisogno di ricordarmi di ripristinare lo stdout originale.

Relazionato:

• Reindirizzare lo stdout a un file in Python?
• Gestione delle eccezioni : articolo interessante sulla gestione delle eccezioni in Python, rispetto a C ++

modificare

So che posso avvolgere target, definire funzioni separate o utilizzare il gestore di contesto . Cerco una soluzione semplice, elegante, idiomatica che non richieda più di 5 righe

Solo pensando fuori dagli schemi qui, che ne dici di un open()metodo personalizzato ?

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()

Usalo in questo modo:

# 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)

Attenersi al codice corrente. È semplice e puoi dire esattamente cosa sta facendo semplicemente guardandolo.

Un altro modo sarebbe con un inline if:

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

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

Ma non è molto più breve di quello che hai e sembra probabilmente peggio.

Potresti anche renderlo non sys.stdoutchiudibile, ma non sembra troppo pitonico:

sys.stdout.close = lambda: None

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

Perché LBYL quando puoi EAFP?

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

Perché riscriverlo per usare il blocco with/ in asmodo uniforme quando devi farlo funzionare in modo contorto? Aggiungerai più linee e ridurrai le prestazioni.

Un'altra possibile soluzione: non cercare di evitare il metodo di uscita del gestore di contesto, basta duplicare lo stdout.

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

Un miglioramento della risposta di 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

Ciò consente l'IO binario e passa eventuali argomenti estranei a openif filenameè effettivamente un nome di file.

Vorrei anche optare per una semplice funzione wrapper, che può essere piuttosto semplice se puoi ignorare la modalità (e di conseguenza stdin vs. stdout), ad esempio:

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 stiamo entrando in guerre a una linea, ecco:

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

Mi piace l'esempio originale di Jacob fintanto che il contesto è scritto solo in un punto. Sarebbe un problema se dovessi riaprire il file per molte scritture. Penso che prenderei la decisione solo una volta all'inizio dello script e lascerei che il sistema chiudesse il file all'uscita:

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

Potresti includere il tuo gestore di uscita se ritieni che sia più ordinato

import atexit

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

atexit(cleanup_output)

Se proprio devi insistere su qualcosa di più "elegante", cioè una battuta:

>>> 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.txtappare e contiene il testo foo.

Che ne dici di aprire un nuovo fd per sys.stdout? In questo modo non avrai problemi a chiuderlo:

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)

Leggero miglioramento in alcuni casi.

import contextlib
import sys

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

Solo due righe extra se stai usando Python 3.3 o versioni successive: una riga per l'extra importe una riga per il stack.enter_context.

Se va bene che sys.stdoutè chiuso dopo il withcorpo, puoi anche usare modelli come questo:

# Use stdout when target is "-"
with open(target, "w") if target != "-" else sys.stdout as f:
    f.write("hello world")

# Use stdout when target is falsy (None, empty string, ...)
with open(target, "w") if target else sys.stdout as f:
    f.write("hello world")

o anche più in generale:

with target if isinstance(target, io.IOBase) else open(target, "w") as f:
    f.write("hello world")