Queens on a chess board

The below pygame program allows the user to place any number of queens on a chess board, and visualize what arrangement of queens will threaten which squares of the chess board.

Here is how it goes (source of arrangements used in the gif here):

The code:

import pygame

pygame.init()
wn = pygame.display.set_mode((600, 600))

class Chess:
    def __init__(self, x, y, w, h, queens):
        self.x = x
        self.y = y
        self.queens = queens
        self.cells = [[pygame.Rect(x + i * w, y + j * h, w, h) for i in range(8)] for j in range(8)]

    def get_pos(self, queen):
        return (queen.rect.x - self.x) // queen.w, (queen.rect.y - self.y) // queen.h
    
    def diag_pos(self, x, y):
        m = min(x, y)
        x, y = x - m, y - m
        while x < len(self.cells[0]) and y < len(self.cells):
            yield self.cells[y][x]
            x += 1
            y += 1

    def diag_neg(self, x, y):
        m = min(len(self.cells[0]) - 1 - x, y)
        x, y = x + m, y - m
        while x + 1 and y < len(self.cells):
            yield self.cells[y][x]
            x -= 1
            y += 1
            
    def vert(self, x, y):
        for r in self.cells:
            yield r[x]

    def hori(self, x, y):
        for cell in self.cells[y]:
            yield cell
            
    def threatened(self, queen):
        x, y = self.get_pos(queen)
        if -1 < x < len(self.cells[0]) and -1 < y < len(self.cells):
            for func in [self.diag_pos, self.diag_neg, self.vert, self.hori]:
                for cell in func(x, y):
                    yield cell

    def draw(self):
        for i, row in enumerate(self.cells):
            for j, col in enumerate(row):
                pygame.draw.rect(wn, [(255, 255, 255), (0, 0, 0)][abs(i - j) % 2], col)

    def draw_threats(self):
        for queen in self.queens:
            for cell in self.threatened(queen):
                if cell not in self.queens:
                    pygame.draw.rect(wn, (35, 35, 35), cell)

    def draw_queens(self):
        for queen in self.queens:
            queen.draw()


class Queen:
    def __init__(self, x, y, w, h):
        self.rect = pygame.Rect(x, y, w, h)
        self.dragging = False
        self.x = x
        self.y = y
        self.w = w
        self.h = h

    def clicked(self, pos):
        return self.rect.collidepoint(pos)

    def drag(self, x, y):
        self.x += x
        self.y += y
        self.rect.x = round(self.x / self.w) * self.w
        self.rect.y = round(self.y / self.h) * self.h

    def draw(self):
        pygame.draw.rect(wn, (255, 0, 0), self.rect)

chess = Chess(100, 100, 50, 50, [Queen(100, 50, 50, 50),
                                 Queen(150, 50, 50, 50),
                                 Queen(200, 50, 50, 50)])

running = True
while running:
    wn.fill((123, 123, 123))
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            for queen in chess.queens:
                if queen.clicked(event.pos):
                    queen.dragging = True
        elif event.type == pygame.MOUSEMOTION:
            for queen in chess.queens:
                if queen.dragging:
                    queen.drag(*event.rel)
        elif event.type == pygame.MOUSEBUTTONUP:
            for queen in chess.queens:
                queen.dragging = False                

    chess.draw()
    chess.draw_threats()
    chess.draw_queens()
    pygame.display.update()

This part of the code allows you to adjust the number of queens on the board:

chess = Chess(100, 100, 50, 50, [Queen(100, 50, 50, 50),
                                 Queen(150, 50, 50, 50),
                                 Queen(200, 50, 50, 50)])

where you can customize each number.

For the Chess class:

• First argument of Chess is the x coordinate of the top-left corner of the chess board
• Second argument of Chess is the y coordinate of the top-left corner of the chess board
• Third argument of Chess is the width of each square of the chess board
• Forth argument of Chess is the height of each square of the chess board
• Fifth argument of Chess is the list of queens that will be occupying the chess board

For the Queen class:

• First argument of Queen is the x coordinate of the top-left corner of the queen
• Second argument of Queen is the y coordinate of the top-left corner of the queen
• Third argument of Queen is the width of the queen
• Forth argument of Queen is the height of the queen