Jak ładnie obsługiwać zarówno `with open (…)`, jak i `sys.stdout`?

Często muszę wyprowadzać dane do pliku lub, jeśli plik nie jest określony, do standardowego wyjścia. Używam następującego fragmentu:

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

Chciałbym go przepisać i jednolicie obsługiwać oba cele.

W idealnym przypadku byłoby to:

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

ale to nie zadziała dobrze, ponieważ sys.stdout jest zamykany przy opuszczaniu withbloku, a ja tego nie chcę. Ja też nie chcę

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

ponieważ musiałbym pamiętać o przywróceniu oryginalnego standardowego wyjścia.

Związane z:

• Przekierować standardowe wyjście do pliku w Pythonie?
• Handling Exceptions - ciekawy artykuł o obsłudze wyjątków w Pythonie w porównaniu do C ++

Edytować

Wiem, że potrafię zawijać target, definiować oddzielną funkcję lub używać menedżera kontekstu . Szukam prostego, eleganckiego, idiomatycznego rozwiązania, które nie wymagałoby więcej niż 5 linek

Po prostu myślisz nieszablonowo, co powiesz na niestandardową open()metodę?

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

Użyj tego w ten sposób:

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

Trzymaj się obecnego kodu. To proste i możesz dokładnie stwierdzić , co robi, po prostu na to spojrzawszy.

Innym sposobem byłoby użycie inline if:

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

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

Ale to nie jest dużo krótsze niż to, co masz i wygląda prawdopodobnie gorzej.

Możesz również sprawić sys.stdout, że zamkniesz się, ale to nie wydaje się zbyt Pythonowe:

sys.stdout.close = lambda: None

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

Dlaczego LBYL, skoro możesz EFF?

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

Po co przepisywać go tak, aby używał jednolicie with/ asblock, kiedy trzeba sprawić, by działał w zawiły sposób? Dodasz więcej linii i zmniejszysz wydajność.

Inne możliwe rozwiązanie: nie próbuj unikać metody wyjścia menedżera kontekstu, po prostu skopiuj standardowe wyjście.

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

Poprawa odpowiedzi Wolpha

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

Pozwala to na binarne We / Wy i przekazywanie ewentualnych dodatkowych argumentów do openif filenamejest rzeczywiście nazwą pliku.

Wybrałbym również prostą funkcję opakowującą, która może być całkiem prosta, jeśli możesz zignorować tryb (iw konsekwencji stdin vs. stdout), na przykład:

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

W porządku, jeśli wchodzimy w wojny jednowierszowe, oto:

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

Podoba mi się oryginalny przykład Jacoba, o ile kontekst jest zapisany tylko w jednym miejscu. Byłoby to problemem, gdybyś ponownie otworzył plik dla wielu zapisów. Myślę, że podjąłbym decyzję raz na początku skryptu i pozwoliłbym systemowi zamknąć plik przy wyjściu:

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

Możesz dołączyć własny program obsługi wyjścia, jeśli uważasz, że jest bardziej uporządkowany

import atexit

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

atexit(cleanup_output)

Jeśli naprawdę musisz nalegać na coś bardziej „eleganckiego”, tj. Jednolinijkowy:

>>> 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.txtpojawia się i zawiera tekst foo.

Co powiesz na otwarcie nowego fd dla sys.stdout? W ten sposób nie będziesz mieć żadnych problemów z zamknięciem go:

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)

Nieznaczna poprawa w niektórych przypadkach.

import contextlib
import sys

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

Tylko dwie dodatkowe linie, jeśli używasz Pythona 3.3 lub nowszego: jedna linia dla dodatkowej importi jedna dla stack.enter_context.

Jeśli jest w porządku, że sys.stdoutjest zamykany po withciele, możesz również użyć takich wzorów:

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

lub jeszcze bardziej ogólnie:

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