How do I optimize and make my code look easier

Question

I'm a begginer to python and I have started using pygame to make a bezier curve program where you can move the points around and it will show the curve. While making it, I thought that it could be heavily optimized since there are a lot of nested if statements and I would like to have more readable code as well as it is kind of a mess.

import sys, pygame, time
pygame.init()

def set_pixel(surface, x, y, r, g, b):
    surface.set_at((x, y), (r, g, b))

def bezier_curve(p1, p2, p3, p4, t):
    x = (1-t)**3*p1[0] + 3*(1-t)**2*t*p2[0] + 3*(1-t)*t**2*p3[0] + t**3*p4[0]
    y = (1-t)**3*p1[1] + 3*(1-t)**2*t*p2[1] + 3*(1-t)*t**2*p3[1] + t**3*p4[1] # some gnarly math from http://en.wikipedia.org/wiki/B%C3%A9zier_curve into python syntax

    return [x, y]

def length(p1, p2):
    return ((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)**0.5

def main():
    size = width, height = 800, 600
    lerp_amount = 0.001

    black = 0, 0, 0
    white = 255, 255, 255
    brighter_blue = 0, 200, 200

    point_circle_size = 7

    p1 = (point_circle_size, height - point_circle_size)
    p2 = (point_circle_size, point_circle_size)
    p3 = (width - point_circle_size, point_circle_size)
    p4 = (width - point_circle_size, height - point_circle_size)

    mouseX = 0
    mouseY = 0

    screen = pygame.display.set_mode(size)

    while True:
        mouseX, mouseY = pygame.mouse.get_pos()
        for event in pygame.event.get():
            if event.type == pygame.QUIT: sys.exit()

        if pygame.mouse.get_pressed()[0]:
            if length(p1, [mouseX, mouseY]) < point_circle_size:
                if (mouseX, mouseY) != p2 or (mouseX, mouseY) != p3 or (mouseX, mouseY) != p4:
                    p1 = (mouseX, mouseY)
            if length(p2, [mouseX, mouseY]) < point_circle_size:
                if (mouseX, mouseY) != p1 or (mouseX, mouseY) != p3 or (mouseX, mouseY) != p4:
                    p2 = (mouseX, mouseY)
            if length(p3, [mouseX, mouseY]) < point_circle_size:
                if (mouseX, mouseY) != p1 or (mouseX, mouseY) != p2 or (mouseX, mouseY) != p4:
                    p3 = (mouseX, mouseY)
            if length(p4, [mouseX, mouseY]) < point_circle_size:
                if (mouseX, mouseY) != p1 or (mouseX, mouseY) != p2 or (mouseX, mouseY) != p3:
                    p4 = (mouseX, mouseY)

        draw_list = []
        lerp_t = 0
        continue_lerp = True
        while continue_lerp:
            lerp_t += lerp_amount
            draw_list.insert(0, bezier_curve(p1, p2, p3, p4, lerp_t))

            if lerp_t > 1:
                continue_lerp = False

        screen.fill(black)
        continue_draw_list = True
        draw_list_index = 0
        while continue_draw_list:
            try:
                pygame.draw.line(screen, brighter_blue, draw_list[draw_list_index], draw_list[draw_list_index + 1], 3)
            except:
                continue_draw_list = False
            draw_list_index += 1

        pygame.draw.circle(screen, white, p1, point_circle_size, 0)
        pygame.draw.circle(screen, white, p2, point_circle_size, 0)
        pygame.draw.circle(screen, white, p3, point_circle_size, 0)
        pygame.draw.circle(screen, white, p4, point_circle_size, 0)

        pygame.draw.line(screen, white, p1, p2, 1)
        pygame.draw.line(screen, white, p3, p4, 1)

        pygame.display.flip()

if __name__ == "__main__":
    main()

All help is appreciated!

Answer 1

Firstly, welcome to python! I also started wth the language a few months ago and love it, particularly if you try to do things in a pythonic way - it just feels right. You're obviously not new to programming so here a few more advanced ideas for clean, idiomatic programming:

One variable for one entity

Where you have variables like p1, p2, ... they are usually all aspects of the same thing or entity. Python gives you some great ways to treat this as a single variable. The most basic is a list, a dict would give you a chance to name the individual entries - at the cost of having them returned in a guaranteed sorted order.

p1 = (point_circle_size, height - point_circle_size)
p2 = (point_circle_size, point_circle_size)
p3 = (width - point_circle_size, point_circle_size)
p4 = (width - point_circle_size, height - point_circle_size)

Can then become

points = [
    (point_circle_size, height - point_circle_size),
    (point_circle_size, point_circle_size),
    (width - point_circle_size, point_circle_size),
    (width - point_circle_size, height - point_circle_size),
    ]

You'll then need to remember that zero-indexing means p1 is now points[0]

Use Python's built in iteration capabilities

Python has awesome iteration support - it's one of the first pythonic things I picked up. For could really be called for each. So now you have a points you can update

pygame.draw.circle(screen, white, p1, point_circle_size, 0)
pygame.draw.circle(screen, white, p2, point_circle_size, 0)
pygame.draw.circle(screen, white, p3, point_circle_size, 0)
pygame.draw.circle(screen, white, p4, point_circle_size, 0)

to

for point in points:
    pygame.draw.circle(screen, white, point, point_circle_size, 0)

Other great things to look at in general are enumerate, which I'll use in an example later to simplify your nested ifs and zip

Get to know and use the standard libraries

There are a few libraries which offer stuff that is not built-in but is "almost". Two key ones to learn about first are collections and itertools. (You'll find the documentation with the central python docs here: https://docs.python.org/3/library/collections.html and here: https://docs.python.org/3/library/itertools.html)

From collections there are two which could be useful for you: namedtuple and deque (double-ended Queue).

namedtuple lets you assign names to the indices of a tuple. This could be useful for your points (it's actually the example given int he docstring for namedtuple!):

points = [
    (point_circle_size, height - point_circle_size),
    (point_circle_size, point_circle_size),
    (width - point_circle_size, point_circle_size),
    (width - point_circle_size, height - point_circle_size),
    ]

can become:

from collections import namedtuple
...

...
coord = namedtuple("coord", [x,y])
...

...
points = [
    coord(point_circle_size, height - point_circle_size),
    coord(point_circle_size, point_circle_size),
    coord(width - point_circle_size, point_circle_size),
    coord(width - point_circle_size, height - point_circle_size),
    ]

then you can reference them using normal indexing or also as points[0].x and points[0].y

So for example the bezier_curve calculation would start:

x = (1-t)**3*p[0].x + 3*(1-t)**2*t*p[1].x + ...

Deque is what to use if you find yourself using .insert(0, ...). Tuples and lists are optimised as FILO - so .append and .pop work great but .insert(0, ...) is slow as the system needs to rearrange everything in memory. A deque is optimised for access at both ends.

I'd suggest using a deque for draw_list:

from collections import deque
...

...
draw_list = deque()
...

...
    draw_list.appendleft(bezier_curve(p1, p2, p3, p4, lerp_t))
...

Extract functions to make your code easier to read

This is valid in any language, I find it very valid in python as good python code reads like english sentences.

All your nested if statements:

if length(p1, [mouseX, mouseY]) < point_circle_size:
    if (mouseX, mouseY) != p2 or (mouseX, mouseY) != p3 or (mouseX, mouseY) != p4:
        p1 = (mouseX, mouseY)
if length(p2, [mouseX, mouseY]) < point_circle_size:
    if (mouseX, mouseY) != p1 or (mouseX, mouseY) != p3 or (mouseX, mouseY) != p4:
        p2 = (mouseX, mouseY)

have a general form and so can be extracted to a function with a well-chosen name (you can probably do better than me at naming it and the variables):

def near(mouse, coords, i, fuzz):
    return ((length(coords[i], mouse) < fuzz) and (mouse = coords[i] or mouse not in coords)

which lets you simplify the original code to

for i, c in enumerate(points):
    if near(mouse, c, i, point_circle_size):
        c = mouse

This is assuming you have used the tip above to put all your points into one list and noticed that mouseX, mouseY could simply be mouse i.e.:

mouse = pygame.mouse.get_pos()

or (if you have implemented coords as a namedtuple):

mouse = coords(pygame.mouse.get_pos())

Uses classes where they make sense

Having done all that, take a look at the idea of using a class to let you group functionality and variables together in a meaningful way. Your code would lend itself to mainly being part of a class bezier_curve(object) which has attributes such as points, internal attributes such as _draw_list and methods such as calc_curve(self, t), lerp(self) and update_points(self, mouse). If that is enough of a pointer to follow then great. If you'd like a few more specific tips, just comment.

Overall

I'm a fan of clean code, TDD and idiomatic programming. Python is, IMHO, a great language to practice all three to create expressive code that reads like a set of instructions where each line is an english sentence. In general, don't worry too much about trying to plan all these kind of ideas in to your coding from the start:

1. Start by writing code in some way that works, however ugly
2. Make sure you are creating and extending test cases up front, just-in-time as you go - that will already guide you towards creating smaller functions and methods that do one thing each and mean that you can refactor afterwards without any fear - just hit the go button on your tests after every small change.
3. Then "do a Raymond": hit the table and think "There must be a better way"(*), refactor, refactor, refactor until you can read your code easily and it looks nice

(*) - search youtube for videos from Raymond Hettinger if you are interested in learning python this way - he's great to listen to and imparts amazing advice. A good one to start with would be: https://www.youtube.com/watch?v=OSGv2VnC0go

Most importantly - have fun!

Answer 2

I doubt if the bezier_curve function will benefit in terms of [measurable] performance but the code could be reduced to:

def bezier_curve(p1, p2, p3, p4, t):
    mt = 1-t
    return [mt**3*pa + 3*mt**2*t*pb + 3*mt*t**2*pc + t**3*pd for pa, pb, pc, pd in zip(p1, p2, p3, p4)]

Answer 3

I would say this

pygame.draw.circle(screen, white, p1, point_circle_size, 0)
pygame.draw.circle(screen, white, p2, point_circle_size, 0)
pygame.draw.circle(screen, white, p3, point_circle_size, 0)
pygame.draw.circle(screen, white, p4, point_circle_size, 0)

Can be written as

for i in {p1, p2, p3, p4}:
    pygame.draw.circle(screen, white, i, point_circle_size, 0)