fix(livecodebench): support sys.stdout.buffer in output capture

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):

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函数"""