base_imports.py

BASE_IMPORTS = """
from string import *
from re import *
from datetime import *
from collections import *
from heapq import *
from bisect import *
from copy import *
from math import *
from random import *
from statistics import *
from itertools import *
from functools import *
from operator import *
from io import *
from sys import *
from json import *
from builtins import *
from typing import *

import string
import re
import datetime
import collections
import heapq
import bisect
import copy
import math
import random
import statistics
import itertools
import functools
import operator
import io
import sys
import json

sys.setrecursionlimit(6*10**5)
"""

checker.py

class Checker:

    @staticmethod
    def compare_strings(output : str, expected_output : str):
        return output == expected_output
    
    @staticmethod
    def compare_ints(output : int, expected_output : int):
        return output == expected_output
    
    @staticmethod
    def compare_floats(output : float, expected_output : float):
        return output == expected_output

    @staticmethod
    def compare_bools(output : bool, expected_output : bool):
        return output == expected_output

    @staticmethod
    def compare_lists(output : list, expected_output : list):
        if len(output) != len(expected_output):
            return False
        for i in range(len(output)):
            if not Checker.compare(output[i], expected_output[i]):
                return False
        return True

    @staticmethod
    def compare_dicts(output : dict, expected_output : dict):
        if len(output) != len(expected_output):
            return False
        for key in output:
            if key not in expected_output:
                return False
            if not Checker.compare(output[key], expected_output[key]):
                return False
        return True

    @staticmethod
    def compare(output, expected_output):
        if isinstance(expected_output, str):
            return Checker.compare_strings(str(output), expected_output)
        elif isinstance(expected_output, int):
            return Checker.compare_ints(int(output), expected_output)
        elif isinstance(expected_output, float):
            return Checker.compare_floats(float(output), expected_output)
        elif isinstance(expected_output, bool):
            return Checker.compare_bools(bool(output), expected_output)
        elif isinstance(expected_output, (list, tuple)):
            return Checker.compare_lists(list(output), list(expected_output))
        elif isinstance(expected_output, dict):
            return Checker.compare_dicts(dict(output), dict(expected_output))
        else:
            raise Exception(f"Unsupported type: {type(expected_output)}")

leetcode_checker.py

import os
import io
import json
import tqdm
import signal
import tempfile
import platform
import traceback
import contextlib
import faulthandler
from enum import Enum
import multiprocessing as mp

from checker import Checker
from base_imports import BASE_IMPORTS

from test_cases import Test


class TimeoutException(Exception):
    pass


class ResultEnum(Enum):
    ACCEPTED = "Accepted"
    COMPILE_ERROR = "Compile Error"
    WRONG_ANSWER = "Wrong Answer"
    RUNTIME_ERROR = "Runtime Error"
    TIME_LIMIT_EXCEEDED = "Time Limit Exceeded"
    UNKNOWN_ERROR = "Unknown Error"


def truncate_str(s, max_len=1000):
    if len(s) > max_len:
        s = s[:max_len] + "...."
    return s


def truncate_io(io, max_len=1000):
    if isinstance(io, (list, tuple)):
        io = "\n".join([truncate_str(str(i), max_len // len(io)) for i in io])
    else:
        assert isinstance(io, str), io
        io = truncate_str(io, max_len)
    return io


@contextlib.contextmanager
def time_limit(seconds):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")

    signal.setitimer(signal.ITIMER_REAL, seconds)
    signal.signal(signal.SIGALRM, signal_handler)
    try:
        yield
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)


@contextlib.contextmanager
def chdir(root):
    if root == ".":
        yield
        return
    cwd = os.getcwd()
    os.chdir(root)
    try:
        yield
    except BaseException as exc:
        raise exc
    finally:
        os.chdir(cwd)


@contextlib.contextmanager
def create_tempdir():
    with tempfile.TemporaryDirectory() as dirname:
        with chdir(dirname):
            yield dirname


def reliability_guard(maximum_memory_bytes=8 * 1024 * 1024 * 1024):
    """
    This disables various destructive functions and prevents the generated code
    from interfering with the test (e.g. fork bomb, killing other processes,
    removing filesystem files, etc.)
    WARNING
    This function is NOT a security sandbox. Untrusted code, including, model-
    generated code, should not be blindly executed outside of one. See the
    Codex paper for more information about OpenAI's code sandbox, and proceed
    with caution.
    """

    if maximum_memory_bytes is not None:
        import resource

        resource.setrlimit(
            resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)
        )
        resource.setrlimit(
            resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)
        )
        if not platform.uname().system == "Darwin":
            resource.setrlimit(
                resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)
            )

    faulthandler.disable()

    import builtins

    builtins.exit = None
    builtins.quit = None

    import os

    os.environ["OMP_NUM_THREADS"] = "1"

    os.kill = None
    os.system = None
    os.putenv = None
    os.remove = None
    os.removedirs = None
    os.rmdir = None
    os.fchdir = None
    os.setuid = None
    os.fork = None
    os.forkpty = None
    os.killpg = None
    os.rename = None
    os.renames = None
    os.truncate = None
    os.replace = None
    os.unlink = None
    os.fchmod = None
    os.fchown = None
    os.chmod = None
    os.chown = None
    os.chroot = None
    os.fchdir = None
    os.lchflags = None
    os.lchmod = None
    os.lchown = None
    os.getcwd = None
    os.chdir = None

    import shutil

    shutil.rmtree = None
    shutil.move = None
    shutil.chown = None

    import subprocess

    subprocess.Popen = None  # type: ignore

    __builtins__["help"] = None

    import sys

    sys.modules["ipdb"] = None
    sys.modules["joblib"] = None
    sys.modules["resource"] = None
    sys.modules["psutil"] = None
    sys.modules["tkinter"] = None


def check_code(
    code: str, tests: list[Test], func_name: str, early_stop=True, timeout=5
):
    manager = mp.Manager()
    result = manager.list()

    p = mp.Process(target=check_code_process, args=(code, tests, func_name, result))

    p.start()
    p.join(timeout=(timeout + 2) * len(tests) + 5)
    if p.is_alive():
        p.kill()

    if not result:
        print("TLE HERE :\??", result)
        result = [
            (
                False,
                [ResultEnum.TIME_LIMIT_EXCEEDED],
                {
                    "last_testcase": "TLE should not happen?",
                    "expected_output": "TLE should not happen?",
                },
            )
        ]

    return result[0]


def check_code_process(
    code: str,
    tests: list[Test],
    func_name: str,
    result: list,
    early_stop=True,
    timeout=5,
):
    # print(code, func_name)
    with io.StringIO() as buf, contextlib.redirect_stdout(buf):
        with create_tempdir():

            # These system calls are needed when cleaning up tempdir.
            import os
            import shutil

            rmtree = shutil.rmtree
            rmdir = os.rmdir
            chdir = os.chdir
            cwd = os.getcwd

            reliability_guard()

            global_result, results, metadata = check_code_helper(
                code, tests, func_name, early_stop=early_stop, timeout=timeout
            )

            # Needed for cleaning up.
            shutil.rmtree = rmtree
            os.rmdir = rmdir
            os.chdir = chdir
            os.getcwd = cwd

    # print("!!!", global_result, results, metadata)
    result.append((global_result, results, metadata))


def check_code_helper(
    code: str, tests: list[Test], func_name: str, early_stop=True, timeout=5
):
    results = []
    exec_dict = {}
    exec(BASE_IMPORTS, exec_dict)
    try:
        exec(code, exec_dict)
    except Exception as e:
        return (
            False,
            [ResultEnum.COMPILE_ERROR],
            {
                "error": truncate_io(str(e)),
                "traceback": truncate_io(traceback.format_exc()),
            },
        )

    metadata = None

    ordered_tests = tests
    for test_idx, test in enumerate(ordered_tests):
        inputs = test.input
        inputs = [json.loads(input) for input in inputs.split("\n")]
        expected_output = json.loads(test.output)
        try:
            new_code_string = f"MY_OUTPUT_RESULT = Solution().{func_name}(*MY_INPUTS)"
            exec_dict["MY_INPUTS"] = inputs
            with time_limit(timeout):
                exec(new_code_string, exec_dict)
            generated_output = exec_dict["MY_OUTPUT_RESULT"]
            if not Checker.compare(generated_output, expected_output):
                results.append(ResultEnum.WRONG_ANSWER)
                metadata = {
                    "last_testcase": truncate_io(inputs),
                    "expected_output": truncate_io(json.dumps(expected_output)),
                    "code_output": truncate_io(json.dumps(generated_output)),
                }
                if early_stop:
                    break
            else:
                results.append(ResultEnum.ACCEPTED)
        except TimeoutException:
            results.append(ResultEnum.TIME_LIMIT_EXCEEDED)
            metadata = {
                "last_testcase": truncate_io(inputs),
                "expected_output": truncate_io(json.dumps(expected_output)),
            }
            if early_stop:
                break
        except Exception as e:
            results.append(ResultEnum.RUNTIME_ERROR)
            metadata = {
                "last_testcase": truncate_io(inputs),
                "expected_output": truncate_io(json.dumps(expected_output)),
                "error": (traceback.format_exc()),
                "error-short": (repr(e)),
            }
            if early_stop:
                break

    global_result = all([result == ResultEnum.ACCEPTED for result in results])
    return global_result, results, metadata


def run_code_only(code: str, test: Test, func_name: str, timeout=5):
    manager = mp.Manager()
    result = manager.list()

    p = mp.Process(
        target=run_code_only_helper, args=(code, test, func_name, result, timeout)
    )

    p.start()
    p.join(timeout=(timeout + 2))
    if p.is_alive():
        p.kill()
        result = [ResultEnum.TIME_LIMIT_EXCEEDED]

    return result[0]


def run_code_only_helper(
    code: str, test: Test, func_name: str, result: list, timeout=5
):
    exec_dict = {}
    exec(BASE_IMPORTS, exec_dict)
    try:
        with io.StringIO() as buf, contextlib.redirect_stdout(buf):
            exec(code, exec_dict)
    except Exception as e:
        result.append(ResultEnum.COMPILE_ERROR)
        print(result)
        return
    inputs = test.input
    try:
        new_code_string = f"MY_OUTPUT_RESULT = Solution().{func_name}(*MY_INPUTS)"
        exec_dict["MY_INPUTS"] = inputs
        with time_limit(timeout):
            with io.StringIO() as buf, contextlib.redirect_stdout(buf):
                exec(new_code_string, exec_dict)
        generated_output = exec_dict["MY_OUTPUT_RESULT"]
        result.append(generated_output)
        return
    except TimeoutException:
        result.append(ResultEnum.TIME_LIMIT_EXCEEDED)
        # print(result)
        return
    except:
        try:
            new_code_string = f"MY_OUTPUT_RESULT = Solution().{func_name}(MY_INPUTS)"
            exec_dict["MY_INPUTS"] = inputs
            with time_limit(timeout):
                with io.StringIO() as buf, contextlib.redirect_stdout(buf):
                    exec(new_code_string, exec_dict)
            generated_output = exec_dict["MY_OUTPUT_RESULT"]
            result.append(generated_output)
            return
        except TimeoutException:
            result.append(ResultEnum.TIME_LIMIT_EXCEEDED)
            return
        except Exception as e:
            # print(repr(e))
            result.append(ResultEnum.RUNTIME_ERROR)
            return

test_cases.py

from enum import Enum
from dataclasses import dataclass

class TestType(Enum):
    STDIN = "stdin"
    FUNCTIONAL = "functional"


@dataclass
class Test:
    input: str
    output: str
    testtype: TestType

    @classmethod
    def from_leetcode_test(cls, test: dict):
        return cls(
            input=test["inputs"], output=test["outputs"], testtype=TestType.FUNCTIONAL
        )

    @classmethod
    def from_codeforces_test(cls, test: dict):
        return cls(input=test["input"], output=test["answer"], testtype=TestType.STDIN)

    def __hash__(self):
        return hash((str(self.input), str(self.output), self.testtype.value))