Paint program made in PyGame

Question

I am making a paint program in Python using Pygame. Right now the program is very laggy, and I wanted to get my code reviewed incase the reason is in the way I have written the code.

EDIT: I was mostly able to fix the performance issue by changing the way layers work. The issue was too much blitting and I fixed it by showing everything on the active layer and only blitting that. It would still be useful if someone reviewed the code.

(Code below is the updated version)

import pygame as pg
import pygame.gfxdraw
import sys,os,math

pg.init()
pg.event.set_allowed([pg.QUIT,pg.MOUSEMOTION,pg.MOUSEBUTTONDOWN,pg.KEYDOWN])

# Settings:
clock = pygame.time.Clock()
fps = 120
font = pg.font.SysFont('consolas',20)
screenres = (500,500)
realres = (screenres[0]*1.2,screenres[1]*1.2)

updated = False
dirtyrects = []

# Colors | R | G | B | A |
clear =  (  0,  0,  0,  0)
white =  (255,255,255)
gray =   (150,150,150)
black =  (  0,  0,  0)
red =    (255,  0,  0)
orange = (255,125,  0)
yellow = (255,255,  0)
green =  (  0,225,  0)
blue =   (  0,  0,255)
purple = (150,  0,150)

colors = [black,white,red,orange,yellow,green,blue,purple]

numkey = [
    pg.K_1,
    pg.K_2,
    pg.K_3,
    pg.K_4,
    pg.K_5,
    pg.K_6,
    pg.K_7,
    pg.K_8
]

# Surfaces:
window = pg.display.set_mode(screenres,pg.DOUBLEBUF)
window.fill(white)
canvas = pg.Surface((realres[0],realres[1]*0.84)).convert_alpha()
canvas.fill(white)
latest1 = canvas.copy()
latest2 = canvas.copy()
latest3 = canvas.copy()
latest4 = canvas.copy()
latest5 = canvas.copy()
layers = [latest1,latest2,latest3,latest4,latest5]
for layer in layers:
    layer.fill(clear)
overlay = pg.Surface(screenres).convert_alpha()

# Rects:
realrect = pg.Rect(0,0,realres[0],int(realres[1]*0.84))
screenrect = pg.Rect(0,0,screenres[0],int(screenres[1]*0.84))
toolbar = pg.Rect(0,420,500,80)

r = 25
clr = black
startpoint = None
endpoint = None
ongoing = False
undone = 0
maxundone = 0
holdingclick = False

def button(color,rect):
    global clr,holdingclick
    if 0 <= rect <= 9:
        rect = pg.Rect(48*rect+12,446,44,44)
        if pg.mouse.get_pressed()[0] and rect.collidepoint(mousepos) and not holdingclick:
            clr = color
            dirtyrects.append(toolbar)
        if clr == color:
            pg.draw.rect(overlay,color,rect)
            pg.draw.rect(overlay,black,rect,3)
        else:
            pg.draw.rect(overlay,color,(rect[0]+4,rect[1]+4,rect[2]-8,rect[3]-8))
            pg.draw.rect(overlay,black,(rect[0]+4,rect[1]+4,rect[2]-8,rect[3]-8),3)

def drawline():
    global startpoint,endpoint,start
    if startpoint == None:
        startpoint = x,y
    endpoint = x,y
    if r > 1:
        if startpoint != endpoint:
            dx,dy = endpoint[0]-startpoint[0],endpoint[1]-startpoint[1]
            angle = math.atan2(-dy,dx)%(2*math.pi)
            dx,dy = math.sin(angle)*(r*0.999),math.cos(angle)*(r*0.999)
            a = startpoint[0]+dx,startpoint[1]+dy
            b = startpoint[0]-dx,startpoint[1]-dy
            c = endpoint[0]-dx,endpoint[1]-dy
            d = endpoint[0]+dx,endpoint[1]+dy
            pointlist = [a,b,c,d]
            pg.draw.polygon(latest1,clr,pointlist)
        pg.draw.circle(latest1,clr,(x,y),r)
    else:
        pg.draw.line(latest1,clr,startpoint,endpoint,r)
    startpoint = x,y

def shiftdown():
    for layer in reversed(layers):
        if layer == latest5:
            canvas.blit(latest5,(0,0))
        else:
            layers[layers.index(layer)+1].blit(layer,(0,0))

def shiftup():
    for layer in layers:
        if layer == latest5:
            layer.fill(clear)
        else:
            layer.fill(clear)
            layer.blit(layers[layers.index(layer)+1],(0,0))

# Drawing static parts of overlay:
overlay.fill(clear)
pg.draw.rect(overlay,gray,toolbar)
pg.draw.rect(overlay,black,toolbar,3)

# Drawing number indicators for colors:
for color in colors:
    text = font.render(str(colors.index(color)+1),True,black)
    overlay.blit(text,(48*colors.index(color)+28,424))

overlaybg = overlay.copy()

while True:
    for event in pg.event.get():
        if event.type == pg.QUIT or pg.key.get_pressed()[pg.K_ESCAPE]:
            pg.quit()
            sys.exit()

        if event.type == pg.MOUSEMOTION:
            mousepos = pg.mouse.get_pos()
            x = int(mousepos[0]*(realres[0]/screenres[0]))
            y = int(mousepos[1]*(realres[1]/screenres[1]))
            holdingclick = True
            if screenrect.collidepoint(mousepos):
                dirtyrects.append(screenrect)

        if event.type == pg.MOUSEBUTTONDOWN:
            holdingclick = False
            if screenrect.collidepoint(mousepos):
                dirtyrects.append(screenrect)

            # Changing brush size:
            if event.button == 4 and r < 100:
                r += 1
                dirtyrects.append(screenrect)
            elif event.button == 5 and r > 2:
                r -= 1
                dirtyrects.append(screenrect)

        if event.type == pg.KEYDOWN:

            if event.key in numkey:
                clr = colors[numkey.index(event.key)]
                dirtyrects.append(toolbar)

            # Emptying canvas:
            if event.key == pg.K_e:
                canvas.fill(white)
                latest5.fill(clear)
                latest4.fill(clear)
                latest3.fill(clear)
                latest2.fill(clear)
                latest1.fill(clear)
                undone = 0
                maxundone = 0
                dirtyrects.append(screenrect)

            # Undoing and redoing:
            if event.key == pg.K_u and undone < maxundone:
                undone += 1
                dirtyrects.append(screenrect)
            if event.key == pg.K_i and undone > 0:
                undone -= 1
                dirtyrects.append(screenrect)

    # Painting:
    if pg.mouse.get_pressed()[0] and screenrect.collidepoint(mousepos):
        if not ongoing:
            while undone > 0:
                shiftup()
                undone -= 1
                maxundone -= 1
            shiftdown()
        drawline()
        ongoing = True
    else:
        startpoint = None
        if ongoing:
            if maxundone < 5:
                maxundone += 1
            ongoing = False

    if screenrect in dirtyrects:

        # Drawing canvas:
        window.fill(white)
        for layer in layers:
            if layers.index(layer) == undone:
                window.blit(pg.transform.smoothscale(layer,(screenrect[2],screenrect[3])),screenrect)

        # Drawing overlay:
        overlay.fill(clear)
        if r > 1:
            pg.gfxdraw.aacircle(overlay,mousepos[0],mousepos[1],int(r*screenres[0]/realres[0]),gray)
    overlay.blit(overlaybg,screenrect)
    for color in colors:
        button(color,colors.index(color))
    window.blit(overlay,screenrect)

    pg.display.set_caption('Draw   |   FPS: ' + str(int(clock.get_fps())))
    clock.tick(fps)

    # Updating display:
    if not updated:
        pg.display.update()
        updated = True
    pg.display.update(dirtyrects)
    dirtyrects.clear()

Answer 1

The are two big problems with the presented code: globals and incorrect use of functions.

Try to avoid globals by any means. The code that uses them becomes more difficult to debug. See here for discussion: Why is Global State so Evil? For example, in the function called button it's like you have two interfaces, one with passed parameters and one with globals.

The functions don't follow Single responsibility principle. In the button function too many things happen: changing color, appending dirty rectangle to a list and drawing rectangle. Try not to put too much logic in a single function. Just an example of how a function should look like, when you draw buttons for selecting colors you could use a function like this:

def toolbar_element_offset(count: int,
                           *,
                           first_offset: int,
                           shift: int = 48) -> int:
    """Distance of element on toolbar from the left border"""
    return first_offset + count * shift

and then call it with different first_offset values for selected or unselected buttons to get their coordinates. Moreover, it could be used to print number labels as well.

Though, I understand that with pygame it will be challenging to stick to a functional programming approach. I'm not a pygame expert (in fact I know almost nothing about it), but I think it's worth to try the OOP approach. Probably, you would want to use Sprites. My advice is to go through examples presented in the tutorials around the web and see how experts organize their code with pygame.

---

Another problem that will definitely cause you problems is Magic numbers. For example, if you would want to change the number of canvas copies from 5 to something else, you would have to change it in two places. Make them constants instead, something like:

CANVAS_BACKUPS_COUNT = 5

Also, as I understand some magic numbers depend on each other, like screen resolution, size of the canvas and their ratio: 500 * 0.84 = 80. You have all three in the code. Make two of them constants, and calculate the third one.

---

Try to follow PEP 8 style guide. Things that I noticed you violate are: imports, blank lines, spaces after commas, equations aligning, comparing with None (instead of if startpoint == None: should be if startpoint is None:), and 79 characters per line. Also, make constants UPPER_CASE. And follow the conventions for variables names. screenres -> screen_resolution, holdingclick -> holding_click, etc. Some names though could be improved, like clr -> current_color. And remove unnecessary comments like # Rects:. It's pretty clear from the code that those are Rects.

Looks like there is no need to keep both layers and latests. Just one list should be enough. Also, instead of writing:

latest1 = canvas.copy()
latest2 = canvas.copy()
latest3 = canvas.copy()
latest4 = canvas.copy()
latest5 = canvas.copy()
layers = [latest1,latest2,latest3,latest4,latest5]

you could write:

layers = [canvas.copy() for _ in range(CANVAS_BACKUPS_COUNT)]

---

This looks inefficient:

for color in colors:
    text = font.render(str(colors.index(color)+1),True,black)

Instead of searching the color in the list of colors, you could simply iterate over range:

for index in range(len(colors)):
    text = font.render(str(index + 1), True, BLACK)

---

This:

if event.button == 4 and r < 100:
    r += 1
elif event.button == 5 and r > 2:
    r -= 1

is unclear for me. To understand what those 4 and 5 mean I had to go to documentation. What you could do is to use enum.IntEnum. And it would be something like this:

from enum import IntEnum
...
class MouseWheel(IntEnum):
    UP = 4
    DOWN = 5
...
if event.button == MouseWheel.UP and brush_size < max_size:
    brush_size += size_delta
elif event.button == MouseWheel.DOWN and brush_size > min_size:
    brush_size -= size_delta

---

Finally, you use lots of tuples for mouse position or sizes of rectangles. I think it's worth to check out pg.Vector2.

Answer 2

Sorry I don't have time rn to properly look all through your code for you but one quick thing I noticed is that it would be much nicer to put your colours in a dictionary. It could be implemented like:

colours = {"black": (  0,  0,  0),
           "white": (255,255,255),
           ...}

And then to get the colours you can just call colours["black"] rather than black which ties them all up together in one data structure which makes your code easier to understand and neater. Also, if you want to use your colors list which you currently have to iterate through (I don't know if you do), then you can use the new colours dictionary with .index()