写一个俄罗斯方块的代码WASD操作要求用Python

以下是一个基本的俄罗斯方块游戏代码，使用WASD操作，用Python实现：

import random

# 定义方块的形状
shapes = [
    [(1, 1, 1), (0, 1, 0)],
    [(0, 2, 2), (2, 2, 0)],
    [(3, 3, 0), (0, 3, 3)],
    [(4, 0, 0), (4, 4, 4)],
    [(0, 0, 5), (5, 5, 5)],
    [(6, 6, 6, 6)],
    [(7, 7), (7, 7)]
]

# 定义方块的颜色
colors = [
    (0, 0, 0),
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
    (0, 255, 255),
    (255, 255, 255)
]

class Block:
    def __init__(self, x, y, shape):
        self.x = x
        self.y = y
        self.shape = shape

    def move(self, dx, dy):
        self.x += dx
        self.y += dy

    def rotate(self):
        self.shape = list(zip(*self.shape[::-1]))

    def get_block(self):
        return [[self.x + j, self.y + i, self.shape[i][j]] for i in range(len(self.shape)) for j in range(len(self.shape[0])) if self.shape[i][j] != 0]

class Tetris:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.board = [[0 for _ in range(width)] for _ in range(height)]
        self.current_block = None

    def new_block(self):
        shape = random.choice(shapes)
        x = self.width // 2 - len(shape[0]) // 2
        y = 0
        self.current_block = Block(x, y, shape)

    def move_block(self, dx, dy):
        if self.current_block:
            if self.is_valid_move(dx, dy):
                self.current_block.move(dx, dy)
                return True
        return False

    def rotate_block(self):
        if self.current_block:
            self.current_block.rotate()
            if not self.is_valid_move(0, 0):
                self.current_block.rotate()
                self.current_block.rotate()
                self.current_block.rotate()
                return False
            return True
        return False

    def drop_block(self):
        if self.current_block:
            while self.is_valid_move(0, 1):
                self.current_block.move(0, 1)
            return True
        return False

    def is_valid_move(self, dx, dy):
        block = self.current_block.get_block()
        for x, y, _ in block:
            if x + dx < 0 or x + dx >= self.width or y + dy >= self.height:
                return False
            if y + dy >= 0 and self.board[y+dy][x+dx] != 0:
                return False
        return True

    def add_block_to_board(self):
        block = self.current_block.get_block()
        for x, y, c in block:
            self.board[y][x] = c
        self.current_block = None

    def remove_full_rows(self):
        new_board = []
        full_rows = 0
        for row in self.board:
            if 0 in row:
                new_board.append(row)
            else:
                full_rows += 1
                new_board.insert(0, [0 for _ in range(self.width)])
        self.board = new_board
        return full_rows

    def get_board(self):
        if self.current_block:
            block = self.current_block.get_block()
        else:
            block = []
        board = [[0 for _ in range(self.width)] for _ in range(self.height)]
        for y in range(self.height):
            for x in range(self.width):
                if self.board[y][x] != 0:
                    board[y][x] = self.board[y][x]
                elif any([xx == x and yy == y for xx, yy, _ in block]):
                    board[y][x] = colors[self.current_block.shape[0][0]]
        return board

# 游戏主循环
def main_loop():
    tetris = Tetris(10, 20)

    clock = pygame.time.Clock()
    fps = 30
    drop_speed = 1
    drop_event = pygame.USEREVENT + 1
    pygame.time.set_timer(drop_event, 1000 // drop_speed)

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    return
                if event.key == pygame.K_LEFT:
                    tetris.move_block(-1, 0)
                if event.key == pygame.K_RIGHT:
                    tetris.move_block(1, 0)
                if event.key == pygame.K_DOWN:
                    tetris.move_block(0, 1)
                if event.key == pygame.K_UP:
                    tetris.rotate_block()
                if event.key == pygame.K_SPACE:
                    tetris.drop_block()
            if event.type == drop_event:
                if not tetris.move_block(0, 1):
                    tetris.add_block_to_board()
                    full_rows = tetris.remove_full_rows()
                    if full_rows > 0:
                        print('Removed {} row(s)'.format(full_rows))
                    tetris.new_block()
                    if not tetris.is_valid_move(0, 0):
                        print('Game Over!')
                        pygame.quit()
                        return

        board = tetris.get_board()
        screen.fill((0, 0, 0))
        for y in range(len(board)):
            for x in range(len(board[y])):
                pygame.draw.rect(screen, colors[board[y][x]], pygame.Rect(x*30, y*30, 30, 30), 0)
        pygame.display.flip()
        clock.tick(fps)

# 初始化pygame
pygame.init()

# 设置窗口大小和标题
screen_width = 300
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Tetris')

# 启动游戏主循环
main_loop()

# 退出pygame
pygame.quit()