Archipelago/worlds/apquest/game/generate_math_problem.py

import random
from collections.abc import Callable
from enum import Enum
from operator import add, mul, sub, truediv
from typing import NamedTuple

_random = random.Random()

class NumberChoiceConstraints(NamedTuple):
    num_1_min: int
    num_1_max: int
    num_2_min: int
    num_2_max: int
    commutative: bool
    operator: Callable[[int, int], int | float]


class MathProblemType(Enum):
    PLUS = 1
    MINUS = 2
    TIMES = 3
    DIVIDE = 4


MATH_PROBLEM_CONSTRAINTS = {
    MathProblemType.PLUS: NumberChoiceConstraints(1, 99, 1, 99, True, add),
    MathProblemType.MINUS: NumberChoiceConstraints(2, 99, 1, 99, False, sub),
    MathProblemType.TIMES: NumberChoiceConstraints(2, 10, 2, 50, True, mul),
    MathProblemType.DIVIDE: NumberChoiceConstraints(4, 99, 2, 50, False, truediv),
}


class MathProblem(NamedTuple):
    problem_type: MathProblemType
    num_1: int
    num_2: int
    result: int


def generate_math_problem(random_object: random.Random = _random) -> MathProblem:
    problem_type: MathProblemType = random_object.choice(list(MathProblemType))
    number_choice_constraints = MATH_PROBLEM_CONSTRAINTS[problem_type]

    for _ in range(10000):
        num_1 = random.randint(number_choice_constraints.num_1_min, number_choice_constraints.num_1_max)
        num_2 = random.randint(number_choice_constraints.num_2_min, number_choice_constraints.num_2_max)

        result = number_choice_constraints.operator(num_1, num_2)

        result_int = int(result)
        if not result_int == result:
            continue

        if result_int < 2 or result_int > 99:
            continue

        if number_choice_constraints.commutative:
            if random.randint(0, 1):
                num_1, num_2 = num_2, num_1

        return MathProblem(problem_type, num_1, num_2, result_int)

    return MathProblem(MathProblemType.PLUS, 1, 1, 2)