A basic tictactoe game

tictactoe.py

from abc import ABC
from enum import Enum
from itertools import product
from typing import Optional

import readchar
from readchar.key import UP, DOWN, LEFT, RIGHT, ENTER

"""
Small TicTacToe game.

How to run:
```
python -m venv venv
source venv/bin/activate
pip install readchar
python tictactoe.py
```

And you can use arrow keys / ENTER to select your moves.

Possible improvements:
* Add a non-human player, with some kind of optimal strategy
* Add some tests
* Refactoring: separate the screen display from the game logic 
"""


# --- Utils ---


class ANSI:
    """Small utility class to format text in a terminal, using ANSI
    escape sequences

    Ref: https://gist.github.com/fnky/458719343aabd01cfb17a3a4f7296797
    """

    @staticmethod
    def underline(txt: str) -> str:
        return f'\033[4m{txt}\033[24m'

    @staticmethod
    def erase_screen() -> str:
        return '\033[2J'


# --- Grid ---

class CellType(Enum):
    empty = " "
    x = "x"
    o = "o"

    def __str__(self):
        return self.value


def _grid_to_str(grid: list[list]) -> str:
    rows = []
    for i in range(Grid.SIZE):
        rows.append("|".join(f" {grid[i][j]} " for j in range(Grid.SIZE)))
        if i < Grid.SIZE - 1:
            rows.append("+".join("---" for _ in range(Grid.SIZE)))
    return "\n".join(rows)


class Grid:
    SIZE = 3

    def __init__(self):
        self.grid = [
            [CellType.empty for _ in range(self.SIZE)] for _ in range(self.SIZE)
        ]

    def play(self, move: CellType, coord: tuple[int, int]):
        """
        Register a new move on the grid. Note that the move must be valid
        (i.e. the target cell must be empty).

        Args:
            move:
            coord:
        """
        i, j = coord
        if not self.is_cell_empty(coord):
            raise ValueError(
                f"Move on coordinate {coord} not allowed "
                f'(already filled with: "{self.grid[i][j]})"'
            )
        self.grid[i][j] = move

    def is_cell_empty(self, coord: tuple[int, int]) -> bool:
        """
        Check if a cell is empty and can be played

        Args:
            coord:

        Returns:
            True if the cell is empty, else False
        """
        i, j = coord
        if not (0 <= i < self.SIZE and 0 <= j < self.SIZE):
            raise ValueError(f"Coordinates {i, j} out of the grid")
        return self.grid[i][j] == CellType.empty

    def is_full(self) -> bool:
        """

        Returns:
            True if the grid is full and no-one can play anymore
        """
        num_empty = sum(
            self.is_cell_empty(coord) for coord in product(range(self.SIZE), repeat=2)
        )
        if num_empty > 0:
            return False
        return True

    def get_winner(self) -> Optional[CellType]:
        """
        Returns:
            The CellType of the winner if a player has won, else None
        """
        # All the index sequences corresponding to winning situations
        seqs = (
            # all columns
            *(tuple((i, j) for i in range(self.SIZE)) for j in range(self.SIZE)),
            # all rows
            *(tuple((i, j) for j in range(self.SIZE)) for i in range(self.SIZE)),
            # top-left to bottom-right diagonal
            tuple((i, i) for i in range(self.SIZE)),
            # bottom-left to top-right diagonal
            tuple((self.SIZE - (i + 1), i) for i in range(self.SIZE)),
        )
        for seq in seqs:
            if all(self.grid[i][j] == CellType.x for i, j in seq):
                return CellType.x
            if all(self.grid[i][j] == CellType.o for i, j in seq):
                return CellType.o
        return None

    def display(self, title: str = '', cursor: Optional[tuple[int, int]] = None):
        print(ANSI.erase_screen())
        print(title)
        grid = [[str(x) for x in row] for row in self.grid]
        if cursor:
            i, j = cursor
            grid[i][j] = ANSI.underline(grid[i][j])
        print(_grid_to_str(grid))


# --- Player ---

class Player(ABC):
    def get_move_coord(self, grid: Grid) -> tuple[int, int]:
        """
        Ask the player which move they want to play.

        Args:
            grid: the current grid

        Returns:
            The coordinates they chose to play on the grid
        """


def _update_cursor(cursor: tuple[int, int], move: tuple[int, int]) -> tuple[int, int]:
    """Try to move the cursor inside the grid. The move is applied only if it is
    valid (i.e. the cursor stays inside the grid)

    Returns:
        The new cursor position
    """
    new_cursor = cursor[0] + move[0], cursor[1] + move[1]
    if all(0 <= i < Grid.SIZE for i in new_cursor):
        return new_cursor
    return cursor


class HumanPlayer(Player):
    def __init__(self, player_name: str):
        self.name = player_name

    def get_move_coord(self, grid: Grid) -> tuple[int, int]:
        """Let the player choose a valid coordinate on the grid.
        To move the cursor: arrow keys
        To choose an (empty) coordinate: ENTER
        """
        # (0, 0) is top left
        # (2, 2) is bottom right
        _moves = {
            UP: (-1, 0),
            DOWN: (1, 0),
            LEFT: (0, -1),
            RIGHT: (0, 1),
        }

        cursor = (1, 1)  # starting in the middle of the grid
        grid.display(title=f'Player turn: {self.name}', cursor=cursor)
        while True:
            key = readchar.readkey()
            if key in (UP, DOWN, LEFT, RIGHT):
                # Get a new cursor position and update grid
                cursor = _update_cursor(cursor, _moves[key])
                grid.display(title=f'Player turn: {self.name}', cursor=cursor)
            if key == ENTER and grid.is_cell_empty(cursor):
                return cursor


# --- Game ---

class Game:
    def __init__(self):
        print("Game starting...\n")
        first_name = input("Please enter 1st player name: ")
        second_name = input("Please enter 2nd player name: ")
        self.players = (
            (CellType.x, HumanPlayer(first_name)),
            (CellType.o, HumanPlayer(second_name)),
        )
        self.grid = Grid()

    def play(self):
        """Run the game loop"""
        i = 0
        over = False
        self.grid.display()
        while not over:
            move, player = self.players[i % 2]
            coord = player.get_move_coord(self.grid)
            self.grid.play(move, coord)
            self.grid.display()
            over = self.check_game_status()
            i += 1

    def check_game_status(self) -> bool:
        """
        Returns:
            True if the game is now over, False otherwise
        """
        if (cell_type := self.grid.get_winner()) is not None:
            winner = next(p for ct, p in self.players if ct == cell_type)
            print(f"player {winner.name} won!")
            return True
        if self.grid.is_full():
            print("grid is full, that's a draw")
            return True
        return False


def main():
    game = Game()
    game.play()


if __name__ == "__main__":
    main()