Як добре обробляти і `з відкритим (...)`, і `sys.stdout`?

Часто мені потрібно виводити дані або у файл, або, якщо файл не вказаний, у stdout. Я використовую такий фрагмент:

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

Я хотів би переписати його і однаково обробляти обидві цілі.

В ідеальному випадку це буде:

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

але це не буде працювати добре, тому що sys.stdout закривається при виході з withблоку, і я цього не хочу. Я ні хочу

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

тому що мені потрібно було б пам'ятати, щоб відновити оригінальний stdout.

Пов’язані:

• Перенаправити stdout на файл у Python?
• Обробка винятків - цікава стаття про обробку винятків у Python порівняно з C ++

Редагувати

Я знаю, що можу обернути target, визначити окрему функцію або використовувати менеджер контексту . Я шукаю просте, елегантне, ідіоматичне рішення, що не вимагає більше 5 рядків

Думаючи тут нестандартно, як щодо власного open()методу?

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

Використовуйте його так:

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

Дотримуйтесь поточного коду. Це просто, і ви можете точно сказати що він робить, просто подивившись на це.

Інший спосіб - вбудований if:

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

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

Але це не набагато коротше, ніж у вас, і виглядає, можливо, гірше.

Ви також можете зробити sys.stdoutзакритим, але це не здається надто пітонічним:

sys.stdout.close = lambda: None

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

Чому LBYL, коли ви можете EAFP?

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

Навіщо переписувати його, щоб використовувати with/ asblock рівномірно, коли вам потрібно змусити його працювати заплутано? Ви додасте більше рядків і знизите продуктивність.

Інше можливе рішення: не намагайтеся уникати методу виходу з менеджера контексту, просто продублюйте stdout.

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

Покращення відповіді Вольфа

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

Це дозволяє двійковому вводу-виводу і передавати можливі сторонні аргументи, openякщо filenameце справді ім'я файлу.

Я б також вибрав просту функцію обгортки, яка може бути досить простою, якщо ви можете ігнорувати режим (і, отже, stdin проти stdout), наприклад:

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

Добре, якщо ми вступаємо у однолінійні війни, ось:

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

Мені подобається оригінальний приклад Якова, якщо контекст написаний лише в одному місці. Було б проблемою, якщо ви в кінцевому підсумку знову відкриєте файл для багатьох записів. Я думаю, я б просто прийняв рішення один раз у верхній частині сценарію і дозволив системі закрити файл при виході:

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

Ви можете включити власний обробник виходу, якщо вважаєте, що він більш охайний

import atexit

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

atexit(cleanup_output)

Якщо ви дійсно повинні наполягати на чомусь більш "елегантному", тобто однокласному:

>>> 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.txtз'являється і містить текст foo.

Як щодо відкриття нового fd для sys.stdout? Таким чином у вас не буде проблем із закриттям:

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)

У деяких випадках незначне поліпшення.

import contextlib
import sys

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

Лише два додаткові рядки, якщо ви використовуєте Python 3.3 або вище: один рядок для додаткового importі один рядок для stack.enter_context.

Якщо це нормально, якщо sys.stdoutзакрито після withтіла, ви також можете використовувати такі моделі:

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

або навіть загальніше:

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