diff --git a/atpbar/stream/output.py b/atpbar/stream/output.py index 61ef5bd..fca90a6 100644 --- a/atpbar/stream/output.py +++ b/atpbar/stream/output.py @@ -18,17 +18,11 @@ def __init__(self, queue: StreamQueue, fd: FD) -> None: self.buffer = '' def write(self, s: str) -> int: - # sys.__stdout__.write(repr(s)) - # sys.__stdout__.write('\n') + if not isinstance(s, str): + # The same error message as `sys.stdout.write()` + raise TypeError(f'write() argument must be str, not {type(s).__name__}') - try: - endswith_n = s.endswith('\n') - except: - self.flush() - self.queue.put((s, self.fd)) - return len(s) - - if endswith_n: + if s.endswith('\n'): self.buffer += s self.flush() return len(s) diff --git a/tests/stream/output/test_stateful.py b/tests/stream/output/test_output_stream.py similarity index 89% rename from tests/stream/output/test_stateful.py rename to tests/stream/output/test_output_stream.py index 5e419e0..fbd7e18 100644 --- a/tests/stream/output/test_stateful.py +++ b/tests/stream/output/test_output_stream.py @@ -1,5 +1,6 @@ from unittest.mock import sentinel +import pytest from hypothesis import given, settings from hypothesis import strategies as st @@ -7,6 +8,12 @@ from tests.stream.st import st_text +def test_type_error() -> None: + stream = OutputStream(Queue(), sentinel.fd) + with pytest.raises(TypeError): + stream.write(123) # type: ignore + + class StatefulTest: def __init__(self, data: st.DataObject) -> None: self.draw = data.draw