diff --git a/ais_bench/benchmark/datasets/livecodebench/testing_util.py b/ais_bench/benchmark/datasets/livecodebench/testing_util.py index c77e6715..e12e2299 100644 --- a/ais_bench/benchmark/datasets/livecodebench/testing_util.py +++ b/ais_bench/benchmark/datasets/livecodebench/testing_util.py @@ -11,6 +11,7 @@ from datetime import datetime from enum import Enum # for capturing the stdout +import io from io import StringIO # used for testing the code that reads from input from unittest.mock import mock_open, patch @@ -63,11 +64,43 @@ 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 _StdoutWithBuffer(StringIO): + """StringIO that also exposes a binary ``.buffer`` view, like real stdout. + + Generated solutions often use ``sys.stdout.buffer.write(...)`` for fast + output. A plain StringIO has no ``buffer`` attribute, so such solutions + used to fail with an AttributeError and were counted as runtime errors. + The bytes written through ``.buffer`` are decoded and forwarded to this + same StringIO, so ``getvalue()`` still returns the complete output. + """ + + class _Buffer: + + def __init__(self, owner): + self._owner = owner + + def write(self, data): + if isinstance(data, (bytes, bytearray)): + data = data.decode() + return self._owner.write(data) + + def writelines(self, lines): + for line in lines: + self.write(line) + + def flush(self): + pass + + def __init__(self): + super().__init__() + self.buffer = _StdoutWithBuffer._Buffer(self) + + class Capturing(list): def __enter__(self): self._stdout = sys.stdout - sys.stdout = self._stringio = StringIO() + sys.stdout = self._stringio = _StdoutWithBuffer() # Make closing the StringIO a no-op self._stringio.close = lambda x: 1 return self @@ -665,6 +698,19 @@ def stripped_string_compare(s1, s2): return s1 == s2 +class _StdinWithBuffer(StringIO): + """StringIO that also exposes a binary ``.buffer`` view, like real stdin. + + Generated solutions often use ``sys.stdin.buffer.read()`` for fast input. + A plain StringIO has no ``buffer`` attribute, so such solutions used to + fail with an AttributeError and were counted as runtime errors. + """ + + def __init__(self, data): + super().__init__(data) + self.buffer = io.BytesIO(data.encode()) + + def call_method(method, inputs): if isinstance(inputs, list): @@ -676,7 +722,7 @@ def call_method(method, inputs): # @patch('builtins.input', side_effect=inputs.split("\n")) @patch('builtins.open', mock_open(read_data=inputs)) - @patch('sys.stdin', StringIO(inputs)) + @patch('sys.stdin', _StdinWithBuffer(inputs)) @patch('sys.stdin.readline', lambda *args: next(inputs_line_iterator)) @patch('sys.stdin.readlines', lambda *args: inputs.split('\n')) @patch('sys.stdin.read', lambda *args: inputs) diff --git a/tests/UT/datasets/livecodebench/test_livecodebench.py b/tests/UT/datasets/livecodebench/test_livecodebench.py index 5ec83b13..c773fe17 100644 --- a/tests/UT/datasets/livecodebench/test_livecodebench.py +++ b/tests/UT/datasets/livecodebench/test_livecodebench.py @@ -427,6 +427,62 @@ def test_load_with_cot(self): self.assertIn('train', result) +class TestStdIOBuffer(unittest.TestCase): + """测试标准输入/输出的二进制 .buffer 视图 + + 竞赛类题解常用 sys.stdin.buffer.read() / sys.stdout.buffer.write() 做快速 + 读写,评测时伪造的 stdin/stdout 必须提供 .buffer,否则这类题解会抛 + AttributeError 被误判为 Runtime Error。 + """ + + @classmethod + def setUpClass(cls): + try: + from ais_bench.benchmark.datasets.livecodebench import testing_util + cls.testing_util = testing_util + except ImportError: + raise unittest.SkipTest('testing_util not available') + + def _run(self, source, inputs): + namespace = {} + exec(compile(source, '', 'exec'), namespace) + with self.testing_util.Capturing() as output: + self.testing_util.call_method(namespace['main'], inputs) + return ''.join(output).strip() + + def test_stdin_buffer_read(self): + """sys.stdin.buffer.read() 应读到测试输入""" + source = ('import sys\n' + 'def main():\n' + ' data = sys.stdin.buffer.read().split()\n' + ' print(sum(map(int, data)))\n') + self.assertEqual(self._run(source, '1 2\n3\n'), '6') + + def test_stdin_buffer_readline(self): + """sys.stdin.buffer.readline() 应按行读到测试输入""" + source = ('import sys\n' + 'def main():\n' + ' n = int(sys.stdin.buffer.readline())\n' + ' print(n * 2)\n') + self.assertEqual(self._run(source, '21\n'), '42') + + def test_stdout_buffer_write(self): + """sys.stdout.buffer.write() 的输出应被 Capturing 收集到""" + source = ('import sys\n' + 'def main():\n' + ' sys.stdout.buffer.write(b"42\\n")\n') + self.assertEqual(self._run(source, ''), '42') + + def test_mixed_text_and_binary_output(self): + """混用 print 和 sys.stdout.buffer.write 时输出顺序应保持一致""" + source = ('import sys\n' + 'def main():\n' + ' print("a")\n' + ' sys.stdout.buffer.write(b"b\\n")\n' + ' print("c")\n') + self.assertEqual(self._run(source, ''), 'a\nb\nc') + + if __name__ == '__main__': unittest.main()