Làm cách nào để xử lý tốt cả `với open (…)` và `sys.stdout`?

Thường thì tôi cần xuất dữ liệu ra tệp hoặc, nếu tệp không được chỉ định, sang stdout. Tôi sử dụng đoạn mã sau:

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

Tôi muốn viết lại nó và xử lý thống nhất cả hai mục tiêu.

Trong trường hợp lý tưởng, nó sẽ là:

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

nhưng điều này sẽ không hoạt động tốt vì sys.stdout bị đóng khi rời khỏi withkhối và tôi không muốn điều đó. Tôi không muốn

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

bởi vì tôi cần phải nhớ khôi phục stdout ban đầu.

Có liên quan:

• Chuyển hướng stdout đến một tệp bằng Python?
• Xử lý ngoại lệ - bài viết thú vị về xử lý ngoại lệ trong Python, so với C ++

Biên tập

Tôi biết rằng tôi có thể bao bọc target, xác định chức năng riêng biệt hoặc sử dụng trình quản lý ngữ cảnh . Tôi tìm kiếm một giải pháp đơn giản, thanh lịch, phù hợp thành ngữ, không yêu cầu nhiều hơn 5 dòng

Chỉ cần suy nghĩ bên ngoài hộp ở đây, làm thế nào về một open()phương pháp tùy chỉnh ?

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

Sử dụng nó như thế này:

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

Gắn bó với mã hiện tại của bạn. Nó đơn giản và bạn có thể biết chính xác nó đang làm gì chỉ bằng cách nhìn lướt qua nó.

Một cách khác sẽ là với nội tuyến if:

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

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

Nhưng nó không ngắn hơn nhiều so với những gì bạn có và nó có vẻ tệ hơn.

Bạn cũng có thể làm cho không thể sys.stdoutmở được, nhưng điều đó có vẻ không quá Pythonic:

sys.stdout.close = lambda: None

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

Tại sao LBYL khi bạn có thể EAFP?

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

Tại sao phải viết lại nó để sử dụng thống nhất with/ asblock khi bạn phải làm cho nó hoạt động một cách phức tạp? Bạn sẽ thêm nhiều dòng và giảm hiệu suất.

Một giải pháp khả thi khác: đừng cố tránh phương thức thoát trình quản lý ngữ cảnh, chỉ cần sao chép stdout.

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

Một cải tiến cho câu trả lời của 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

Điều này cho phép IO nhị phân và chuyển các đối số không liên quan đến opennếu filenamethực sự là tên tệp.

Tôi cũng muốn sử dụng một hàm wrapper đơn giản, có thể khá đơn giản nếu bạn có thể bỏ qua chế độ (và do đó stdin so với stdout), ví dụ:

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

Được rồi, nếu chúng ta đang tham gia vào cuộc chiến một bên, đây là:

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

Tôi thích ví dụ ban đầu của Jacob miễn là ngữ cảnh chỉ được viết ở một nơi. Sẽ là một vấn đề nếu bạn mở lại tệp trong nhiều lần ghi. Tôi nghĩ rằng tôi sẽ chỉ đưa ra quyết định một lần ở đầu tập lệnh và để hệ thống đóng tệp khi thoát:

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

Bạn có thể bao gồm trình xử lý thoát của riêng mình nếu bạn nghĩ rằng nó gọn gàng hơn

import atexit

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

atexit(cleanup_output)

Nếu bạn thực sự phải nhấn mạnh vào một cái gì đó "thanh lịch" hơn, tức là một lớp lót:

>>> 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.txtxuất hiện và chứa văn bản foo.

Làm thế nào về việc mở một fd mới cho sys.stdout? Bằng cách này, bạn sẽ không gặp bất kỳ sự cố nào khi đóng nó:

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)

Cải thiện nhẹ trong một số trường hợp.

import contextlib
import sys

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

Chỉ cần thêm hai dòng nếu bạn đang sử dụng Python 3.3 trở lên: một dòng cho phần bổ sung importvà một dòng cho stack.enter_context.

Nếu nó sys.stdoutđược đóng sau withphần thân, bạn cũng có thể sử dụng các mẫu như sau:

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

hoặc thậm chí tổng quát hơn:

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