Wordle solver

wordle.py

import typing as _t
from collections import defaultdict
from enum import Enum
from random import choice


class TileColor(Enum):
    GRAY = 0
    YELLOW = 1
    GREEN = 2


class Wordle:
    def __init__(self, secret: str, words: _t.Iterator[str]) -> None:
        self.secret = secret.lower()
        self.words = {_ for _ in words if len(_) == len(secret)}

    def submit(self, word: str) -> list[TileColor]:
        assert len(word) == len(self.secret), "word has same length as secret"
        assert word in self.words, "word is in the dictionary"

        ss = set(self.secret)
        hints = []
        for a, b in zip(word.lower(), self.secret):
            if a == b:
                hints.append(TileColor.GREEN)
            elif a in ss:
                hints.append(TileColor.YELLOW)
                ss.remove(a)
            else:
                hints.append(TileColor.GRAY)
        return hints


class WordleSolver:
    def __init__(self, words: list[str]) -> None:
        self.words = words

        self.exclude = set()
        self.correct = {}
        self.misplaced = defaultdict(set)
        self.include = set()

    def guess(self) -> str:
        return choice(self.words)

    def shrink(self, word: str, tiles: list[TileColor]) -> list[str]:
        for i, (c, t) in enumerate(zip(word, tiles)):
            if t is TileColor.GRAY:
                self.exclude.add(c)
            elif t is TileColor.GREEN:
                self.correct[i] = c
            else:
                self.misplaced[i].add(c)
                self.include.add(c)

        feasible = [
            (
                len(
                    {c for i, c in enumerate(word) if i not in self.correct}
                    & self.include
                ),
                word,
            )
            for word in self.words
            if not (
                self.exclude & set(word)
                or not all(word[i] == c for i, c in self.correct.items())
                or any(
                    word[i] in self.misplaced[i]
                    for i in range(len(word))
                    if i not in self.correct
                )
            )
        ]

        hi = max(s for s, _ in feasible)
        self.words = [w for s, w in feasible if s == hi]


if __name__ == "__main__":

    def wordlist():
        return (
            _.lower()
            for _ in (_.strip() for _ in open("/usr/share/dict/words"))
            if all(c.isalpha() for c in _)
        )

    w = Wordle("siege", wordlist())
    ws = WordleSolver([_ for _ in wordlist() if len(_) == 5])

    for i in range(6):
        guess = ws.guess()
        outcome = w.submit(guess)
        if set(outcome) == {TileColor.GREEN}:
            break
        ws.shrink(guess, outcome)
    else:
        raise RuntimeError("I couldn't solve the Wordle :(")
    i += 1
    print(f"{guess.upper()} (guessed after {i} attempt{'s' if i != 1 else ''})")