Skip to content

Commit

Permalink
only write to stdout buffer when buffer object is available
Browse files Browse the repository at this point in the history
  • Loading branch information
thijskramer committed Oct 19, 2021
1 parent bcb88c4 commit 9cde73f
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 1 deletion.
10 changes: 9 additions & 1 deletion src/flake8/formatting/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,15 @@ def _write(self, output: str) -> None:
if self.output_fd is not None:
self.output_fd.write(output + self.newline)
if self.output_fd is None or self.options.tee:
sys.stdout.buffer.write(output.encode() + self.newline.encode())
# sys.stdout might be replaced, and thus could also be
# a file like object like io.StringIO, which do not
# support the 'buffer' attribute.
try:
sys.stdout.buffer.write(
output.encode() + self.newline.encode()
)
except AttributeError:
sys.stdout.write(output + self.newline)

def write(self, line: Optional[str], source: Optional[str]) -> None:
"""Write the line either to the output file or stdout.
Expand Down
27 changes: 27 additions & 0 deletions tests/unit/test_base_formatter.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,33 @@ def test_write_produces_stdout(capsys):
assert capsys.readouterr().out == f"{line}\n{source}\n"


def test_write_without_stdout_buffer():
"""Verify that stdout writes when it has no buffer."""
line = "Something to write"
source = "source"

buffer_side_effect = AttributeError(
"'_io.StringIO' object has no \
attribute 'buffer'"
)
buffer_write_mock = mock.Mock(side_effect=buffer_side_effect)

write_mock = mock.Mock()
mock.patch("flake8.formatting.base.sys.stdout.write", write_mock)

with mock.patch(
"flake8.formatting.base.sys.stdout.buffer.write", buffer_write_mock
):
formatter = base.BaseFormatter(options())
formatter.write(line, source)
assert write_mock.called is True
assert write_mock.call_count == 2
assert write_mock.mock_calls == [
mock.call(line + formatter.newline),
mock.call(source + formatter.newline),
]


class AfterInitFormatter(base.BaseFormatter):
"""Subclass for testing after_init."""

Expand Down

0 comments on commit 9cde73f

Please sign in to comment.