mastermind.py

#!/usr/bin/env python
import random
from enum import Enum
from typing import NamedTuple
from collections import Counter
from dataclasses import dataclass

class Color(Enum):
    RED = 1
    ORANGE = 2
    YELLOW = 3
    GREEN = 4
    BLUE = 5
    PURPLE = 6

    def __repr__(self):
        return str(self)

@dataclass
class Result:
    correct: int = 0
    almost: int = 0

    def __init__(self, target, colors):
        remaining = Counter()
        unmatched = []
        for i in range(4):
            if colors[i] == target[i]:
                self.correct += 1
            else:
                remaining[target[i]] += 1
                unmatched.append(colors[i])
        for color in unmatched:
            if remaining[color] > 0:
                self.almost += 1
                remaining[color] -= 1


class Game:
    def __init__(self, default=None):
        self.target_colors = [random.choice(list(Color)) for i in range(4)]
        if default is not None:
            self.target_colors = default
        self.guesses = []

    def __str__(self):
        return str(self.target_colors)

    def guess(self, colors):
        result = Result(self.target_colors, colors)
        self.guesses.append((colors, result))
        return result