Lightning (Python3/Pygame)

#!/usr/bin/python3
"""
Lightning
by Christopher Chen (christoipher.chen@gmail.com)

Dependencies:
    Python 3
    Pygame

Controls:
    SPACE or Click: create lightning at mouse cursor
    r: create firework at random position
    ESCAPE or q: quit
    F12: take screenshot
    Numpad number keys: create fireworks at coordinates corresponding to the positions of the numpad number keys. Nunmlock must be OFF.

Additional Notes:
    This programm will automatically create a folder called 'Masterpieces' alongside the source py file.
    This is the screenshot directory.
"""



FRAMERATE = 60
FRAMEDELAY = 1/FRAMERATE
key_dict = {}
game_obj_list = []
RANDCOLOR = (255,255,255)



class Streamer:
    """A randomly turning projectile that draws its trail"""
    def __init__(self, x, y, color, speed, lifetime, deviance):
        self.x = x
        self.y = y
        self.color = color
        # speed in pixels/second
        self.speed = speed
        self.angle = random.randint(1,360)
        # self.framerate = framerate
        self.killflag = False
        self.lifetime = lifetime
        self.pos_hist = [(self.x,self.y)]
        self.deviance = deviance

    def update(self):
        self.pos_hist.append((self.x,self.y))
        self.x += self.speed*math.cos(math.radians(self.angle))
        self.y += self.speed*math.sin(math.radians(self.angle))
        self.angle += random.randint(int(-self.deviance/2),int(self.deviance/2))

        # pygame.draw.line(DISPLAYSURF,self.color,(self.prevx,self.prevy),(self.x,self.y),1)
        if self.lifetime <= 0:
            self.killflag = True
        self.lifetime -= 1
        # print("yes")

    def draw(self):
        if FANCY_GRAPHICS:
            pygame.draw.aalines(DISPLAYSURF, self.color, False, self.pos_hist, 1)
        else:
            for index,item in enumerate(self.pos_hist):
                try:
                    pygame.draw.line(DISPLAYSURF,self.color,self.pos_hist[index+1],item,1)
                except IndexError:
                    pass

    def should_die(self):
        return self.killflag

class Firework:
    """A parent for multiple streamers"""

    def __init__(self, n, x, y, color, speed, lifetime, deviance):
        for i in range(n):
            game_obj_list.append(Streamer(x,y,color,speed,lifetime,deviance))
        self.killflag = True

    def should_die(self):
        return True

    def update(self):
        pass

    def draw(self):
        pass

def create_firework_random(x,y):
    color = (random.randint(0,255),random.randint(0,255),random.randint(0,255))
    # color = (random.randint(1,16)*16-1,random.randint(1,16)*16-1,random.randint(1,16)*16-1)
    #                    n,                            x,y,color,speed,                 lifetime,  deviance
    game_obj_list.append(Firework(random.randint(2,15),x,y,color,random.randint(10,80),random.randint(2,10),random.randint(10,180)))

def getkeystate(k):
    try:
        if key_dict[k] == KEYDOWN:
            return True
        elif key_dict[k] == KEYUP:
            return False
    except:
        key_dict.update({k:KEYUP})
        return False

def update_key_dict():
    for event in pygame.event.get():
            #~ print(event)
            if event.type == QUIT:
                pygame.quit()
                sys.exit()

            try:
                key_dict.update({event.key:event.type})
            except AttributeError:
                pass


def save_screenshot():
    screeny_dir = os.path.join(os.curdir,"Masterpieces")
    if not os.path.exists(os.path.join(os.curdir,"Masterpieces")):
        os.makedirs(os.path.join(os.curdir,"Masterpieces"))
    pygame.image.save(DISPLAYSURF,os.path.join(screeny_dir, "Masterpiece " + time.strftime("%Y%m%d %H%M%S")+".png"))

def main():
    while True:
        # get a random color
        RANDCOLOR = (random.randint(0,255),random.randint(0,255),random.randint(0,255))

        #~ handle events and update dictionary of key states
        update_key_dict()

        # get the mouse position
        mouse_pos = {'x':pygame.mouse.get_pos()[0],'y':pygame.mouse.get_pos()[1]}
        # print(mouse_pos)

        # game_obj_list.append(Streamer(mouse_pos['x'],mouse_pos['y'],(255,255,255),10,50,45))
        #create firework
        if getkeystate(K_SPACE) or pygame.mouse.get_pressed()[0] == 1:
            create_firework_random(mouse_pos['x'],mouse_pos['y'])
        #create firewok at random position
        if getkeystate(K_r):
            create_firework_random(random.randint(0,WIDTH-1),random.randint(0,HEIGHT-1))

        #quit
        if getkeystate(K_q) or getkeystate(K_ESCAPE):
            sys.exit()

        #screenshot
        if getkeystate(K_F12):
            save_screenshot()

        for key in numpad_keys:
            if getkeystate(key):
                create_firework_random(WIDTH/2 + (WIDTH/3) * numpad_coords[key][0],HEIGHT/2 - (HEIGHT/3) * numpad_coords[key][1])


        if pygame.mouse.get_pressed()[2] == 0 and not getkeystate(K_LSHIFT):
            DISPLAYSURF.fill((0,0,0))
        time.sleep(FRAMEDELAY)


        # update all items in the main list of game objects
        for index,item in enumerate(game_obj_list):
            if item.should_die():
                del game_obj_list[index]
            item.update()
            item.draw()

        pygame.display.update()

if __name__ == '__main__':
    import random, pygame, math, time, sys, os
    from pygame.locals import *

    pygame.init()
    infoObj = pygame.display.Info()
    WIDTH = int(infoObj.current_w)
    HEIGHT = int(infoObj.current_h)
    SIZE  = (WIDTH,HEIGHT)
    DISPLAYSURF = pygame.display.set_mode((WIDTH,HEIGHT),pygame.FULLSCREEN)

    #numpad_keys = [K_KP1,K_KP2,K_KP3,K_KP4,K_KP5,K_KP6,K_KP7,K_KP8,K_KP9]
    numpad_coords = {K_KP7:(-1,1),K_KP8:(0,1),K_KP9:(1,1),
    K_KP4:(-1,0),K_KP5:(0,0),K_KP6:(1,0),
    K_KP1:(-1,-1),K_KP2:(0,-1),K_KP3:(1,-1)}

    numpad_keys = numpad_coords.keys()

    print("Run with -nogfx to disable fancy graphics")
    if "-nogfx" in sys.argv:
        FANCY_GRAPHICS = False
    else:
        FANCY_GRAPHICS = True
    main()