Tic-tac-toe game

tic_tac_toe.py

# Maciej Goszczycki 2019
# Human is always True. We invert the board when drawing to show the right symbol.
from functools import lru_cache
import random
import re


def colour(colour, text):
    return f"\033[{30 + colour}m{text}\033[0m"


def straight(board, i):
    row = board[i] == board[3 + i] == board[6 + i] is not None
    col = board[i * 3] == board[i * 3 + 1] == board[i * 3 + 2] is not None
    return row or col


def place(board, pos, piece):
    return board[:pos] + (piece,) + board[pos + 1 :]


def tie(board):
    return all(x is not None for x in board)


def won(board):
    diag_lr = board[0] == board[4] == board[8] is not None
    diag_rl = board[2] == board[4] == board[6] is not None
    return diag_lr or diag_rl or any(straight(board, i) for i in range(3))


def optimal(board, turn):
    return {
        p: wins(place(board, p, turn), not turn) for p in range(9) if board[p] is None
    }


@lru_cache(maxsize=None)
def wins(board, turn):
    if won(board):
        return 1 if turn else -1
    if tie(board):
        return 0
    return (min if turn else max)(optimal(board, turn).values())


def ai(board):
    options = [k for k, v in optimal(board, False).items() if v >= 0]
    return place(board, random.choice(options), False) if options else board


def player():
    while True:
        if (choice := input(colour(6, "1-9? ")).strip()) in list("123456789"):
            if board[int(choice) - 1] is None:
                return int(choice) - 1
        print(colour(31, "invalid move. type a single digit"))


def draw(board, start="\n"):
    if human == "x":
        board = [None if x is None else not x for x in board]
    lut = {False: "X", True: "O", None: "."}
    r = "\n".join(" ".join(lut[board[y * 3 + x]] for x in range(3)) for y in range(3))
    print(start + r + "\n")


board = (None,) * 9
while (human := input(colour(6, "x or o? ")).strip().lower()) not in list("xo"):
    print(colour(1, "type 'x' or 'o'"))
if human == "o":
    board = ai(board)
draw(board)
while not tie(board) and not won(board):
    board = place(board, player(), True)
    draw(board)
    board = ai(board)
    draw(board, colour(3, "ai move:\n\n"))
print(colour(1, "You lose") if won(board) else colour(4, "Tie!"))
draw(board)