diff --git a/chromaterm/__main__.py b/chromaterm/__main__.py index 0ea1a34c..5f6c4981 100644 --- a/chromaterm/__main__.py +++ b/chromaterm/__main__.py @@ -1,5 +1,6 @@ """The command line utility for ChromaTerm""" import argparse +import io import os import re import select @@ -208,6 +209,13 @@ def process_input(config, data_fd, forward_fd=None, max_wait=None): max_wait (float): The maximum time to wait with no data on either of the file descriptors. None will block until at least one ready to be read. """ + # Avoid BlockingIOError when output to non-blocking stdout is too fast + try: + os.set_blocking(sys.stdout.fileno(), True) + except io.UnsupportedOperation: + # In case stdout is replaced with StringIO + pass + fds = [data_fd] if forward_fd is None else [data_fd, forward_fd] buffer = '' diff --git a/tests/test__main__.py b/tests/test__main__.py index 01343054..eaacbb5a 100644 --- a/tests/test__main__.py +++ b/tests/test__main__.py @@ -274,6 +274,19 @@ def test_parse_rule_group_out_of_bounds(): assert re.search(msg_re, chromaterm.__main__.parse_rule(rule)) +def test_process_input_blocking_stdout(): + """Ensure that `stdout` is put into a blocking state. Otherwise, it triggers + a `BlockingIOError` if it is not ready to be written to. chromaterm#93.""" + pipe_r, _ = os.pipe() + config = chromaterm.__main__.Config() + + os.set_blocking(sys.stdout.fileno(), False) + assert not os.get_blocking(sys.stdout.fileno()) + + chromaterm.__main__.process_input(config, pipe_r, max_wait=0) + assert os.get_blocking(sys.stdout.fileno()) + + def test_process_input_decode_error(capsys): """Attempt to decode a character that is not UTF-8.""" pipe_r, pipe_w = os.pipe()