diff --git a/ais_bench/benchmark/datasets/livecodebench/testing_util.py b/ais_bench/benchmark/datasets/livecodebench/testing_util.py index c0f62430..b65caded 100644 --- a/ais_bench/benchmark/datasets/livecodebench/testing_util.py +++ b/ais_bench/benchmark/datasets/livecodebench/testing_util.py @@ -11,7 +11,7 @@ from datetime import datetime from enum import Enum # for capturing the stdout -from io import BytesIO, StringIO, TextIOWrapper +from io import BytesIO, TextIOWrapper # used for testing the code that reads from input from unittest.mock import mock_open, patch import inspect @@ -63,13 +63,26 @@ def timeout_handler(signum, frame): # used to capture stdout as a list # from https://stackoverflow.com/a/16571630/6416660 # alternative use redirect_stdout() from contextlib +class _CapturedStdout(TextIOWrapper): + """Capture text and binary writes in their original order.""" + + def __init__(self): + super().__init__(BytesIO(), encoding='utf-8', write_through=True) + + def getvalue(self): + self.flush() + return self.buffer.getvalue().decode('utf-8', errors='replace') + + def close(self): + # Generated solutions may close stdout; keep the captured output readable. + pass + + class Capturing(list): def __enter__(self): self._stdout = sys.stdout - sys.stdout = self._stringio = StringIO() - # Make closing the StringIO a no-op - self._stringio.close = lambda x: 1 + sys.stdout = self._stringio = _CapturedStdout() return self def __exit__(self, *args): diff --git a/tests/UT/datasets/livecodebench/test_testing_util.py b/tests/UT/datasets/livecodebench/test_testing_util.py index a8c4ab40..e5a0d163 100644 --- a/tests/UT/datasets/livecodebench/test_testing_util.py +++ b/tests/UT/datasets/livecodebench/test_testing_util.py @@ -82,6 +82,44 @@ def test_capturing_context_manager(self): self.assertIn('test output', output[0]) self.assertIn('another line', output[0]) + def test_capturing_binary_stdout(self): + with Capturing() as output: + sys.stdout.buffer.write('中文\n'.encode('utf-8')) + + self.assertEqual(output, ['中文\n']) + + def test_capturing_preserves_mixed_write_order(self): + with Capturing() as output: + sys.stdout.write('first ') + sys.stdout.buffer.write(b'second ') + print('third') + + self.assertEqual(output, ['first second third\n']) + + def test_capturing_split_utf8_binary_writes(self): + encoded = '中'.encode('utf-8') + with Capturing() as output: + for byte in encoded: + sys.stdout.buffer.write(bytes([byte])) + + self.assertEqual(output, ['中']) + + def test_capturing_survives_stdout_close(self): + with Capturing() as output: + sys.stdout.buffer.write(b'answer\n') + sys.stdout.close() + + self.assertEqual(output, ['answer\n']) + + def test_call_method_can_write_binary_stdout(self): + def solution(): + sys.stdout.buffer.write(b'42\n') + + with Capturing() as output: + call_method(solution, '') + + self.assertEqual(output, ['42\n']) + class TestTruncatefn(TestingUtilTestBase): """测试truncatefn函数"""