From 8c96cd2c290abebc3f67e9d70416433f90057d00 Mon Sep 17 00:00:00 2001 From: Ke Zheng Date: Wed, 30 Sep 2026 14:21:23 +0800 Subject: [PATCH] Fix binary standard I/O in code-generation judge --- lcb_runner/evaluation/testing_util.py | 118 +++++++++++++++++--------- tests/test_binary_stdio.py | 86 +++++++++++++++++++ 2 files changed, 163 insertions(+), 41 deletions(-) create mode 100644 tests/test_binary_stdio.py diff --git a/lcb_runner/evaluation/testing_util.py b/lcb_runner/evaluation/testing_util.py index ecc89474..644c1bae 100644 --- a/lcb_runner/evaluation/testing_util.py +++ b/lcb_runner/evaluation/testing_util.py @@ -1,4 +1,5 @@ import ast +import codecs import json import sys import faulthandler @@ -12,7 +13,7 @@ import numpy as np -from io import StringIO +from io import BytesIO, TextIOWrapper # used for testing the code that reads from input from unittest.mock import patch, mock_open @@ -56,51 +57,94 @@ 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 _CapturingStream(TextIOWrapper): + def close(self): + # Solutions may close stdout; keep the captured bytes available. + self.flush() + + 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 + self._buffer = BytesIO() + self._stream = _CapturingStream( + self._buffer, + encoding="utf-8", + errors="surrogateescape", + newline="\n", + write_through=True, + ) + sys.stdout = self._stream return self def __exit__(self, *args): - self.append(self._stringio.getvalue()) - del self._stringio # free up some memory - sys.stdout = self._stdout + try: + self._stream.flush() + self.append(self._buffer.getvalue().decode("utf-8", "surrogateescape")) + finally: + sys.stdout = self._stdout + TextIOWrapper.close(self._stream) + del self._stream + del self._buffer -# Custom mock for sys.stdin that supports buffer attribute -class MockStdinWithBuffer: +# Keep text and binary reads on the same cursor without text-wrapper read-ahead. +class MockBuffer(BytesIO): def __init__(self, inputs: str): - self.inputs = inputs - self._stringio = StringIO(inputs) - self.buffer = MockBuffer(inputs) - - def read(self, *args): - return self.inputs - - def readline(self, *args): - return self._stringio.readline(*args) + super().__init__(inputs.encode("utf-8")) - def readlines(self, *args): - return self.inputs.split("\n") - def __getattr__(self, name): - # Delegate other attributes to StringIO - return getattr(self._stringio, name) - - -class MockBuffer: +class MockStdinWithBuffer: def __init__(self, inputs: str): - self.inputs = inputs.encode("utf-8") # Convert to bytes + self.buffer = MockBuffer(inputs) + self.encoding = "utf-8" + self.errors = "strict" + + def _read_chars(self, size: int, stop_at_newline: bool = False): + decoder = codecs.getincrementaldecoder(self.encoding)(self.errors) + chars = [] + while len(chars) < size: + byte = self.buffer.read(1) + if not byte: + break + char = decoder.decode(byte) + if char: + chars.append(char) + if stop_at_newline and char == "\n": + break + return "".join(chars) + + def read(self, size: int = -1): + if size is None or size < 0: + return self.buffer.read().decode(self.encoding, self.errors) + return self._read_chars(size) + + def readline(self, size: int = -1): + if size is None or size < 0: + return self.buffer.readline().decode(self.encoding, self.errors) + return self._read_chars(size, stop_at_newline=True) + + def readlines(self, hint: int = -1): + lines = [] + total = 0 + while line := self.readline(): + lines.append(line) + total += len(line) + if hint > 0 and total >= hint: + break + return lines + + def __iter__(self): + return self - def read(self, *args): - # Return as byte strings that can be split - return self.inputs + def __next__(self): + line = self.readline() + if not line: + raise StopIteration + return line - def readline(self, *args): - return self.inputs.split(b"\n")[0] + b"\n" + def __getattr__(self, name): + return getattr(self.buffer, name) def clean_if_name(code: str) -> str: @@ -156,20 +200,12 @@ def call_method(method, inputs): if isinstance(inputs, list): inputs = "\n".join(inputs) - inputs_line_iterator = iter(inputs.split("\n")) - - # Create custom stdin mock with buffer support mock_stdin = MockStdinWithBuffer(inputs) # sys.setrecursionlimit(10000) - # @patch('builtins.input', side_effect=inputs.split("\n")) @patch("builtins.open", mock_open(read_data=inputs)) - @patch("sys.stdin", mock_stdin) # Use our custom mock instead of StringIO - @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) - # @patch('sys.stdout.write', print) + @patch("sys.stdin", mock_stdin) def _inner_call_method(_method): try: return _method() diff --git a/tests/test_binary_stdio.py b/tests/test_binary_stdio.py new file mode 100644 index 00000000..b8caba69 --- /dev/null +++ b/tests/test_binary_stdio.py @@ -0,0 +1,86 @@ +import signal +import sys +import unittest + +from lcb_runner.evaluation.testing_util import ( + Capturing, + MockBuffer, + MockStdinWithBuffer, + call_method, + grade_stdio, + timeout_handler, +) + + +class BinaryStdioTests(unittest.TestCase): + def test_mock_buffer_advances_for_read_and_readline(self): + stream = MockBuffer("ab\ncd\n") + self.assertEqual(stream.read(1), b"a") + self.assertEqual(stream.readline(), b"b\n") + self.assertEqual(stream.readline(), b"cd\n") + self.assertEqual(stream.readline(), b"") + + def test_text_and_binary_reads_share_a_cursor(self): + stream = MockStdinWithBuffer("first\nsecond\nthird\n") + self.assertEqual(stream.readline(), "first\n") + self.assertEqual(stream.buffer.readline(), b"second\n") + self.assertEqual(stream.readline(), "third\n") + self.assertEqual(stream.buffer.readline(), b"") + + def test_sized_text_reads_count_characters(self): + stream = MockStdinWithBuffer("中A\nend") + self.assertEqual(stream.read(1), "中") + self.assertEqual(stream.readline(2), "A\n") + self.assertEqual(stream.read(), "end") + + def test_readlines_and_iteration_use_remaining_input(self): + stream = MockStdinWithBuffer("a\nb\nc") + self.assertEqual(stream.readline(), "a\n") + self.assertEqual(stream.readlines(), ["b\n", "c"]) + self.assertEqual(list(stream), []) + + def test_call_method_mixes_input_methods(self): + seen = [] + + def solution(): + seen.append(input()) + seen.append(sys.stdin.buffer.readline().decode().strip()) + seen.append(sys.stdin.readline().strip()) + + call_method(solution, "one\ntwo\nthree\n") + self.assertEqual(seen, ["one", "two", "three"]) + + def test_capture_preserves_text_binary_write_order(self): + original_stdout = sys.stdout + with Capturing() as output: + sys.stdout.write("a") + sys.stdout.buffer.write(b"b") + print("c") + sys.stdout.close() + self.assertIs(sys.stdout, original_stdout) + self.assertEqual(output, ["abc\n"]) + + def test_stdio_judge_accepts_repeated_binary_reads(self): + code = ( + "import sys\n" + "a = int(sys.stdin.buffer.readline())\n" + "b = int(sys.stdin.buffer.readline())\n" + "print(a + b)\n" + ) + signal.signal(signal.SIGALRM, timeout_handler) + grades, _ = grade_stdio(code, ["1\n2\n"], ["3\n"], timeout=3) + self.assertEqual(grades, [True]) + + def test_stdio_judge_accepts_binary_output(self): + code = ( + "import sys\n" + "value = sys.stdin.buffer.read().strip()\n" + "sys.stdout.buffer.write(value + b'!\\n')\n" + ) + signal.signal(signal.SIGALRM, timeout_handler) + grades, _ = grade_stdio(code, ["hello\n"], ["hello!\n"], timeout=3) + self.assertEqual(grades, [True]) + + +if __name__ == "__main__": + unittest.main()