Simple rainfall animation in Pygame

Question

I made the following code. It's a simple animation of 'raindrops' falling and splattering once they reach the ground. The images I use can be found here: https://i.sstatic.net/oKOVA.jpg

import pygame, sys
from pygame import locals  
import random
import itertools


def listenToQuit():
    for event in pygame.event.get():
        if event.type == locals.QUIT:
            pygame.quit()
            sys.exit()

def makeNewDrop(display, drops):
    # there isa 15% chance we'll make a new drop, each frame
    x = display.get_width() * random.random() # a random xpostion
    return RainDrop(x, display, drops)        

class RainDrop(pygame.sprite.Sprite):
    img = pygame.image.load('assests/rain_drop.png')

    def __init__(self, x, display, group):
        # call the super constructor, and pass in the 
        # group you've created and it is automatically added to the group every 
        # time you create an instance of this class
        pygame.sprite.Sprite.__init__(self, group) 
        self.y = 0
        self.x = x
        self.fall_step = 5
        self.display = display
        self.die = False
        self. animation = iter([
            pygame.image.load('assests/splatter_1.png'),
            pygame.image.load('assests/splatter_2.png'),
            pygame.image.load('assests/splatter_3.png'),
            pygame.image.load('assests/splatter_4.png')
        ])

    def update(self):
        self.checkIfShouldDie()
        self.fall()
        if self.die:
            try:
                self.img = next(self.animation)
            except StopIteration:
                self.kill()

        self.display.blit(self.img, (self.x,self.y))

    def fall(self):
        self.y = self.y + self.fall_step



    def checkIfShouldDie(self):
        if self.img.get_rect().bottom + self.y >= self.display.get_height():
            self.die = True
            self.y = self.display.get_height() - self.img.get_rect().bottom - 5

if __name__ == '__main__':
    pygame.init()
    drops = pygame.sprite.Group()
    cooldown = 0

    fps = 30
    clock = pygame.time.Clock()

    main_display = pygame.display.set_mode((400,400),0,32)
    pygame.display.set_caption("rain!") 

    while True:
        main_display.fill((255,255,255))

        if random.random() >= 0.85 and cooldown == 0:
            # a 15% change we'll make a new drop, each frame
            # assuming we didn't make a drop in the last few frames
            drop = makeNewDrop(main_display, drops) # automatically added to the drops group
            cooldown = cooldown + 5

        drops.update()

        # reduce cooldown
        cooldown = 0 if cooldown <= 0 else cooldown - 1
        listenToQuit()
        pygame.display.update()
        clock.tick(30)

How can this code be improved? Specifically the part where I iterator over animations and catch the StopIteration error seems very hacky to me.

Answer 1

PEP8

The official Python style guide will suggest that:

• there be two newlines between functions at the global scope, and only one newline between functions in the class scope (i.e. check_if_should_die)
• functions be named in lower_snake_case, i.e. listen_to_quit

Hard exit

Currently, you have a forever-loop that only exits on sys.exit. Instead, simply return a boolean from listenToQuit (which can be called should_quit), and if the return value is true, break out of the loop.

Abbreviated imports

pygame.sprite.Sprite can just be Sprite if you from pygame.sprite import Sprite.

Typo

assests -> assets

Generator

    self. animation = iter([
        pygame.image.load('assests/splatter_1.png'),
        pygame.image.load('assests/splatter_2.png'),
        pygame.image.load('assests/splatter_3.png'),
        pygame.image.load('assests/splatter_4.png')
    ])

can be

self.animation = iter(
   pygame.image.load(f'assets/splatter_{i}.png')
   for i in range(1, 5)
)

Generally I don't think your iter/StopIteration approach is that bad. That said, you can rework it by changing your outer while True into a for drop in drops, where drops is an iterable through four instances of a drop object tied to a specific asset.