Operator Overloading Python Tutorial

Welcome to part 19 of the intermediate Python programming tutorial series. In this tutorial, we are going to introduce the "special" or "magic" methods of Python, and, with that, talk about operator overloading.

Coming back to our Blob World code, I am going to change the code slightly, giving us a BlueBlob, GreenBlob and RedBlob class, all of which inherit from the Blob class. Notice that none of these classes take in a color parameter at all, and the color itself is actually hard-coded. This allows us to move this class to another file, and doesn't require us to pass a color at all. RedBlobs will always be red.

blobworld.py

import pygame
import random
from blob import Blob

STARTING_BLUE_BLOBS = 10
STARTING_RED_BLOBS = 3
STARTING_GREEN_BLOBS = 5

WIDTH = 800
HEIGHT = 600
WHITE = (255, 255, 255)
BLUE = (0, 0, 255)
RED = (255, 0, 0)

game_display = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Blob World")
clock = pygame.time.Clock()

class BlueBlob(Blob):
    
    def __init__(self, x_boundary, y_boundary):
        Blob.__init__(self, (0, 0, 255), x_boundary, y_boundary)

class RedBlob(Blob):
    
    def __init__(self, x_boundary, y_boundary):
        Blob.__init__(self, (255, 0, 0), x_boundary, y_boundary)

class GreenBlob(Blob):
    
    def __init__(self, x_boundary, y_boundary):
        Blob.__init__(self, (0, 255, 0), x_boundary, y_boundary)
        

def draw_environment(blob_list):
    game_display.fill(WHITE)

    for blob_dict in blob_list:
        for blob_id in blob_dict:
            blob = blob_dict[blob_id]
            pygame.draw.circle(game_display, blob.color, [blob.x, blob.y], blob.size)
            blob.move()
            blob.check_bounds()

    pygame.display.update()
    

def main():
    blue_blobs = dict(enumerate([BlueBlob(WIDTH,HEIGHT) for i in range(STARTING_BLUE_BLOBS)]))
    red_blobs = dict(enumerate([RedBlob(WIDTH,HEIGHT) for i in range(STARTING_RED_BLOBS)]))
    green_blobs = dict(enumerate([GreenBlob(WIDTH,HEIGHT) for i in range(STARTING_GREEN_BLOBS)]))
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()

        draw_environment([blue_blobs,red_blobs,green_blobs])
        clock.tick(60)

if __name__ == '__main__':
    main()

blob.py

import random


class Blob:

    def __init__(self, color, x_boundary, y_boundary, size_range=(4,8), movement_range=(-1,2)):
        self.size = random.randrange(size_range[0],size_range[1])
        self.color = color
        self.x_boundary = x_boundary
        self.y_boundary = y_boundary
        self.x = random.randrange(0, self.x_boundary)
        self.y = random.randrange(0, self.y_boundary)
        self.movement_range = movement_range

    def move(self):
        self.move_x = random.randrange(self.movement_range[0],self.movement_range[1])
        self.move_y = random.randrange(self.movement_range[0],self.movement_range[1])
        self.x += self.move_x
        self.y += self.move_y

    def check_bounds(self):
        if self.x < 0: self.x = 0
        elif self.x > self.x_boundary: self.x = self.x_boundary
        
        if self.y < 0: self.y = 0
        elif self.y > self.y_boundary: self.y = self.y_boundary

Now, consider that we want these blobs maybe to start having some sort of interaction with eachother. Let's say that the RedBlobs will harm BlueBlobs, and that GreenBlobs are edible, and beneficial, to BlueBlobs, and maybe add to their size, which can aid them in fighting the RedBlob attacks. How might we handle this in code? Obviously, we need some sort of collision detection, but then, how might we handle the collision itself in the code? Wouldn't it be nice if, in the event of a collision of a BlueBlob and GreenBlob, we could just do a BlueBlob + GreenBlob, just like that in the code, and things work? At the moment, it wont, but we can actually write code to make it this easy! Modifying our BlueBlob class, we can start by doing:

class BlueBlob(Blob):
    
    def __init__(self, x_boundary, y_boundary):
        Blob.__init__(self, (0, 0, 255), x_boundary, y_boundary)

    def __add__(self, other_blob):

Like our "dunder init" method, we can have a "dunder add" (__add__) method. This will let us define what happens when someone actually uses a + operator with our object. This is called "operator overloading." In our case, what should we put here? We obviously must pass some sort of variable that will be added. Python handles this for us in the background, knowing that what comes after the + operator is the next parameter. In this example, that will be a blob class object, and we can check it's color attribute. From here, we can handle it however we see fit. I suggest:

    def __add__(self, other_blob):
        if other_blob.color == (255, 0, 0):
            self.size -= other_blob.size
            other_blob.size -= self.size
            
        elif other_blob.color == (0, 255, 0):
            self.size += other_blob.size
            other_blob.size = 0
            
        elif other_blob.color == (0, 0, 255):
            # for now, nothing. Maybe later it does something more. 
            pass
        else:
            raise Exception('Tried to combine one or multiple blobs of unsupported colors!')

This will just be a nice simple way to handle for this, and we're not using any globals here. Also, note the comments. We're choosing to not do anything when blue-colored Blobs come into contact. Also, when we come into contact with a red blob, note that FIRST, the size of the blue blob will be reduced. Then this new size will be deducted from the red blob. This is different from simply subtracting each other's sizes from the collision, and it might not be what you intend. If none of our conditions are met, we raise an exception, since someone is attempting to use the operator, and there's no handling for it. We'd not want this to pass silently.

Now, we can add blobs together with our custom handling. For example, we could do:

def main():
    blue_blobs = dict(enumerate([BlueBlob(WIDTH,HEIGHT) for i in range(STARTING_BLUE_BLOBS)]))
    red_blobs = dict(enumerate([RedBlob(WIDTH,HEIGHT) for i in range(STARTING_RED_BLOBS)]))
    green_blobs = dict(enumerate([GreenBlob(WIDTH,HEIGHT) for i in range(STARTING_GREEN_BLOBS)]))


    print('Current blue size: {}. Current red size: {}'.format(str(blue_blobs[0].size),
                                                               str(red_blobs[0].size)))

    blue_blobs[0] + red_blobs[0]
    print('Current blue size: {}. Current red size: {}'.format(str(blue_blobs[0].size),
                                                               str(red_blobs[0].size))

Output:

Current blue size: 4. Current red size: 7
Current blue size: -3. Current red size: 10

Alright, so we've got this fancy handling for the + operator, but how will we know when to use it? We need some sort of collision detection, which is what we're going to work on in the next tutorial.

The next tutorial: Detecting Collisions in our Game Python Tutorial