challenge.py

import io
import random
import re
import subprocess
import tempfile
import tokenize
from collections import defaultdict
from dataclasses import dataclass, field
from enum import StrEnum
from pathlib import Path
from typing import ClassVar, Optional, TypeAlias

ROOT_DIR = Path(__file__).parent.parent
PYRIGHT_BASIC_CONFIG = """
# pyright: analyzeUnannotatedFunctions=true
# pyright: strictParameterNoneValue=true
# pyright: disableBytesTypePromotions=false
# pyright: strictListInference=false
# pyright: strictDictionaryInference=false
# pyright: strictSetInference=false
# pyright: deprecateTypingAliases=false
# pyright: enableExperimentalFeatures=false
# pyright: reportMissingImports=error
# pyright: reportUndefinedVariable=error
# pyright: reportGeneralTypeIssues=error
# pyright: reportOptionalSubscript=error
# pyright: reportOptionalMemberAccess=error
# pyright: reportOptionalCall=error
# pyright: reportOptionalIterable=error
# pyright: reportOptionalContextManager=error
# pyright: reportOptionalOperand=error
# pyright: reportTypedDictNotRequiredAccess=error
# pyright: reportPrivateImportUsage=error
# pyright: reportUnboundVariable=error
# pyright: reportUnusedCoroutine=error
"""


class Level(StrEnum):
    BASIC = "basic"
    INTERMEDIATE = "intermediate"
    ADVANCED = "advanced"
    EXTREME = "extreme"

    @classmethod
    def is_valid_level(cls, level: str):
        return level in cls._value2member_map_


ChallengeName: TypeAlias = str


@dataclass(frozen=True)
class ChallengeKey:
    level: Level
    name: ChallengeName

    @classmethod
    def from_str(cls, key: str):
        """Create a key object from a string like "basic-foo"."""
        level, name = key.split("-", maxsplit=1)
        return cls(Level(level), name)


@dataclass
class Challenge:
    """A challenge object.

    :param hints: An optional string of hints, use markdown syntax.
    """

    CODE_SPLITTER: ClassVar[str] = "\n## End of your code ##\n"

    name: ChallengeName
    level: Level
    code: str
    user_code: str = field(default="", init=False)
    test_code: str = field(default="", init=False)
    hints: Optional[str] = None

    def __post_init__(self):
        self.parse_code()

    def parse_code(self):
        self.user_code, _, self.test_code = self.code.partition(self.CODE_SPLITTER)


@dataclass(frozen=True, slots=True)
class TypeCheckResult:
    message: str
    passed: bool
    debug_info: dict = field(default_factory=dict)  # For debugging purposes


class ChallengeManager:
    """The manager for challenges.

    :param root_dir: The root directory that contains the files of challenges.
    """

    def __init__(self, root_dir: Optional[Path] = None):
        if not root_dir:
            root_dir = ROOT_DIR / "challenges"
        self.challenges: dict[ChallengeKey, Challenge] = self._load_challenges(root_dir)
        self.challenges_groupby_level: dict[Level, list[ChallengeName]]
        self.challenges_groupby_level = self._get_challenges_groupby_level()

    def has_challenge(self, key: ChallengeKey) -> bool:
        return key in self.challenges

    def get_challenge(self, key: ChallengeKey) -> Challenge:
        return self.challenges[key]

    @property
    def challenge_count(self) -> int:
        """The count of challenges."""
        return len(self.challenges)

    def run_challenge(self, key: ChallengeKey, user_code: str) -> TypeCheckResult:
        challenge = self.get_challenge(key)
        # Make sure user code ends with a new line to avoid issue #63.
        return self._type_check_with_pyright(user_code + "\n", challenge.test_code)

    def get_random_challenge(self) -> dict[str, str]:
        level = random.choice(list(self.challenges_groupby_level.keys()))
        name = random.choice(self.challenges_groupby_level[level])
        return {"level": level, "name": name}

    @staticmethod
    def _load_challenges(root_dir: Path) -> dict[ChallengeKey, Challenge]:
        challenges = {}
        for challenge_folder in root_dir.iterdir():
            question_source = challenge_folder / "question.py"
            if not question_source.exists():
                continue

            # Try to read the optional hints file
            hints_file = challenge_folder / "hints.md"
            if hints_file.exists():
                hints = hints_file.read_text(encoding="utf-8").strip()
            else:
                hints = None

            key = ChallengeKey.from_str(challenge_folder.name)
            challenges[key] = Challenge(
                name=key.name,
                level=key.level,
                code=question_source.read_text(encoding="utf-8"),
                hints=hints,
            )
        return challenges

    def _get_challenges_groupby_level(self) -> dict[Level, list[ChallengeName]]:
        groups: defaultdict[str, list[ChallengeName]] = defaultdict(list)

        for challenge in self.challenges.values():
            groups[challenge.level].append(challenge.name)

        # Sort challenge by name alphabetically.
        for challenge_names in groups.values():
            challenge_names.sort()

        # Make sure groups are ordered by level (from easy to hard)
        return {level: groups[level] for level in Level}

    EXPECT_ERROR_COMMENT = "expect-type-error"

    # Pyright error messages look like:
    # `<filename>:<line_no>:<col_no> - <error|warning|information>: <message>`
    # Here we only capture the error messages and line numbers
    PYRIGHT_MESSAGE_REGEX = (
        r"^(?:.+?):(?P<line_number>\d+):[\s\-\d]+(?P<message>error:.+)$"
    )

    @staticmethod
    def _partition_test_code(test_code: str):
        """
        Split test code from an optional Pyright configuration block and return the test portion plus the effective Pyright configuration.
        
        Parameters:
            test_code (str): Combined test code that may include a separator line "\n## End of test code ##\n" followed by additional Pyright configuration.
        
        Returns:
            tuple[str, str]: A tuple (test_code, pyright_basic_config) where `test_code` is the portion before the splitter and `pyright_basic_config` is the base PYRIGHT_BASIC_CONFIG optionally extended with any config found after the splitter.
        """
        TEST_SPLITTER = "\n## End of test code ##\n"

        # PYRIGHT_BASIC_CONFIG aim to limit user to modify the config
        test_code, end_test_comment, pyright_config = test_code.partition(TEST_SPLITTER)
        pyright_basic_config = PYRIGHT_BASIC_CONFIG

        # Replace `## End of test code ##` with PYRIGHT_BASIC_CONFIG
        if end_test_comment:
            pyright_basic_config += pyright_config
        return test_code, pyright_basic_config

    @classmethod
    def _type_check_with_pyright(
        cls, user_code: str, test_code: str
    ) -> TypeCheckResult:
        """
        Run Pyright on combined user and test code (including any embedded Pyright config) and report type-check results.
        
        Parameters:
            user_code (str): The user's source code to be type-checked.
            test_code (str): The test suite code (may include an embedded Pyright config region) appended to the user code.
        
        Returns:
            TypeCheckResult: An object containing `message`, a newline-separated report of Pyright errors and a summary line,
            and `passed`, which is `true` if no reported errors (other than those originating from Pyright config) remain.
        """
        test_code, pyright_basic_config = cls._partition_test_code(test_code)

        code = f"{user_code}{test_code}{pyright_basic_config}"
        buffer = io.StringIO(code)

        # This produces a stream of TokenInfos, example:
        # TokenInfo(type=4 (NEWLINE), string='\n', start=(4, 3), end=(4, 4), line='"""\n'),
        # TokenInfo(type=62 (NL), string='\n', start=(5, 0), end=(5, 1), line='\n')
        # See https://docs.python.org/3/library/tokenize.html#tokenize.tokenize for more details
        tokens = list(tokenize.generate_tokens(buffer.readline))

        # Find all lines that are followed by a comment # expect-type-error
        expect_error_line_numbers = [
            token.start[0]
            for token in tokens
            if token.type == tokenize.COMMENT
            and token.string[1:].strip() == cls.EXPECT_ERROR_COMMENT
        ]
        # Tracks whether an expected error has been reported by type checker.
        error_line_seen_in_err_msg: dict[int, bool] = {
            lineno: False for lineno in expect_error_line_numbers
        }

        with tempfile.NamedTemporaryFile(suffix=".py") as temp:
            temp.write(code.encode())
            temp.flush()
            # TODO: switch to json output to simplify output parsing.
            # https://microsoft.github.io/pyright/#/command-line?id=json-output
            raw_result = subprocess.run(
                ["pyright", "--pythonversion", "3.12", temp.name],
                capture_output=True,
                text=True,
            )
            stdout, stderr = raw_result.stdout, raw_result.stderr
            if stderr:
                return TypeCheckResult(message=stderr, passed=False)
        error_lines: list[str] = []

        # Substract lineno in merged code by user_code_line_len, so that the lineno in
        # error message matches those in the test code editor. Fixed #20.
        user_code_lines_len = len(user_code.splitlines())
        for line in stdout.splitlines():
            m = re.match(cls.PYRIGHT_MESSAGE_REGEX, line)
            if m is None:
                continue
            line_number, message = int(m["line_number"]), m["message"]
            if line_number in error_line_seen_in_err_msg:
                # Each reported error should be attached to a specific line,
                # If it is commented with # expect-type-error, let it pass.
                error_line_seen_in_err_msg[line_number] = True
                continue
            # Error could be thrown from user code too, in which case delta shouldn't be applied.
            error_line = f"%s:{message}"

            if line_number <= user_code_lines_len:
                error_lines.append(error_line % line_number)
            elif line_number <= user_code_lines_len + len(test_code.splitlines()):
                error_lines.append(error_line % (line_number - user_code_lines_len))
            else:
                error_lines.append(error_line % "[pyright-config]")

        # If there are any lines that are expected to fail but not reported by pyright,
        # they should be considered as errors.
        for line_number, seen in error_line_seen_in_err_msg.items():
            if not seen:
                error_lines.append(
                    f"{line_number - user_code_lines_len}: error: Expected type error but instead passed"
                )

        # Error for pyright-config will not fail the challenge
        passed = True
        for error_line in error_lines:
            if error_line.startswith("[pyright-config]"):
                continue
            passed = False

        error_lines.append(f"\nFound {len(error_lines)} errors")

        return TypeCheckResult(message="\n".join(error_lines), passed=passed)


challenge_manager = ChallengeManager()