用pygame设计画图

import pygame
import function

pygame.init()
screen = pygame.display.set_mode((800,800))
screen.fill((255,255,255))

state =  function.drawLine(screen)
pygame.display.flip()
while True:
    #只要有事件产生就会进入for循环
    for event in pygame.event.get():
        #根据判断type的值来判断是什么事件产生了
        if event.type == pygame.QUIT:
            exit()
        elif event.type == pygame.KEYDOWN:
            state = function.selectState(event,screen)
        if state == '画直线':
            function.line(event,screen)
        elif state == '画矩形':
            function.rect(event,screen)
        elif state == '画圆':
            function.circle(event,screen)

import pygame
def drawLine(screen):
    pygame.draw.rect(screen,(255,255,255),(0,0,150,60))
    font = pygame.font.Font('./aa.ttf', 50)
    text = font.render('画直线', False, (0, 255, 0))
    screen.blit(text, (0, 0))
    pygame.display.flip()
    return '画直线'
def drawRect(screen):
    pygame.draw.rect(screen, (255, 255, 255), (0, 0, 150, 60))
    font = pygame.font.Font('./aa.ttf', 50)
    text = font.render('画矩形', False, (0, 255, 0))
    screen.blit(text, (0, 0))
    pygame.display.flip()
    return '画矩形'
def drawCircle(screen):
    pygame.draw.rect(screen, (255, 255, 255), (0, 0, 150, 60))
    font = pygame.font.Font('./aa.ttf', 50)
    text = font.render('画圆', False, (0, 255, 0))
    screen.blit(text, (0, 0))
    pygame.display.flip()
    return '画圆'
def line(event,screen):
    if event.type == pygame.MOUSEBUTTONDOWN:
        first_pos = event.pos
        nowlast_pos = first_pos
        while True:
            for event in pygame.event.get():
                if event.type == pygame.MOUSEMOTION:
                    now_pos = event.pos
                    pygame.draw.line(screen, (255, 255, 255), first_pos, nowlast_pos, 5)
                    pygame.display.flip()
                    pygame.draw.line(screen, (0, 0, 0), first_pos, now_pos, 5)
                    nowlast_pos = now_pos
                    pygame.display.flip()
                elif event.type == pygame.MOUSEBUTTONUP:
                    last_pos = event.pos
                    pygame.draw.line(screen, (0, 0, 0), first_pos, last_pos, 5)
                    pygame.display.flip()
                    break
            if event.type == pygame.MOUSEBUTTONUP:
                break
def rect(event,screen):
    if event.type == pygame.MOUSEBUTTONDOWN:
        first_pos = event.pos
        nowlast_pos = first_pos
        while True:
            for event in pygame.event.get():
                if event.type == pygame.MOUSEMOTION:
                    now_pos = event.pos
                    pygame.draw.rect(screen, (255, 255, 255),
                                     (*first_pos, nowlast_pos[0] - first_pos[0], nowlast_pos[1] - first_pos[1]))
                    pygame.display.flip()
                    pygame.draw.rect(screen, (0, 0, 0),
                                     (*first_pos, now_pos[0] - first_pos[0], now_pos[1] - first_pos[1]))
                    nowlast_pos = now_pos
                    pygame.display.flip()
                elif event.type == pygame.MOUSEBUTTONUP:
                    last_pos = event.pos
                    pygame.draw.rect(screen, (0, 0, 0),
                                     (*first_pos, last_pos[0] - first_pos[0], last_pos[1] - first_pos[1]))
                    pygame.display.flip()
                    break
            if event.type == pygame.MOUSEBUTTONUP:
                break
def circle(event,screen):
    if event.type == pygame.MOUSEBUTTONDOWN:
        first_pos = event.pos
        nowlast_pos = first_pos
        while True:
            for event in pygame.event.get():
                if event.type == pygame.MOUSEMOTION:
                    now_pos = event.pos
                    pygame.draw.circle(screen, (255, 255, 255), first_pos, int(
                        ((nowlast_pos[0] - first_pos[0]) ** 2 + (nowlast_pos[1] - first_pos[1]) ** 2) ** 0.5))
                    pygame.display.flip()
                    pygame.draw.circle(screen, (0, 0, 0), first_pos, int(
                        ((now_pos[0] - first_pos[0]) ** 2 + (now_pos[1] - first_pos[1]) ** 2) ** 0.5))
                    nowlast_pos = now_pos
                    pygame.display.flip()
                elif event.type == pygame.MOUSEBUTTONUP:
                    last_pos = event.pos
                    pygame.draw.circle(screen, (0, 0, 0), first_pos, int(
                        ((last_pos[0] - first_pos[0]) ** 2 + (last_pos[1] - first_pos[1]) ** 2) ** 0.5))
                    pygame.display.flip()
                    break
            if event.type == pygame.MOUSEBUTTONUP:
                break

def selectState(event,screen):
    if event.key == 114:
        # 画矩形
        state = drawRect(screen)
        print(event.key)
    elif event.key == 99:
        # 画圆
        state = drawCircle(screen)
        print(event.key)
    elif event.key == 108:
        # h画直线
        state = drawLine(screen)
        print(event.key)
    return state