Variable sized Tic-Tac-Toe end game detection logic in Python

tictactoe.py

from itertools import permutations
import math

X = 'X'
O = 'O'
_ = ' '



def inter_value_differences(values):
    return [a - b for a, b in zip(values[1:], values[:-1])]


def all_values_in_list_equal(values):
    return all(x == values[0] for x in values)


def straight_lines_in_grid(grid_width):
    def points_form_straight_line(dist):
        return dist in [
            1,  # Points align horizontally
            grid_width,  # Points align vertically
            grid_width + 1,  # Points align diagonally down
            grid_width - 1  # Points align diagonally up
        ]

    for points in permutations(range(grid_width * grid_width), grid_width):
        distances_between_points = inter_value_differences(points)

        if not all_values_in_list_equal(distances_between_points):
            continue

        if points_form_straight_line(distances_between_points[0]):
            yield points


def is_end_game_state(board):
    grid_size = int(math.sqrt(len(board)))

    for straight_line in straight_lines_in_grid(grid_size):
        tokens = [board[i] for i in straight_line]
        if all_values_in_list_equal(tokens) and tokens[0] in [X, O]:
            return True

    return False


# Tests

def tests():
    assert is_end_game_state([
        X, _,
        O, X
    ])

    assert not is_end_game_state([
        X, _,
        O, _
    ])

    assert is_end_game_state([
        X, X, _,
        O, X, O,
        X, X, O
    ])

    assert not is_end_game_state([
        X, _, _,
        O, X, O,
        X, X, O
    ])

    assert is_end_game_state([
        X, _, _, O,
        O, X, O, _,
        X, O, O, _,
        O, O, _, X
    ])

    assert not is_end_game_state([
        X, _, _, O,
        O, X, X, _,
        X, O, O, _,
        O, O, _, X
    ])

    assert not is_end_game_state([
        _, _, _, _,
        O, X, X, _,
        X, O, O, _,
        O, O, _, X
    ])


if __name__ == '__main__':
    import cProfile

    cProfile.run('tests()')