paint.py - a pygame-based, minimalist, alternative to Microsoft Paint

paint.py

import os
import time
import pygame
from pygame.locals import *

pygame.init()

scr_w = 500
scr_h = 500

LINE_WIDTH_DEFAULT = 3
LINE_WIDTH_ERASER  = 12
BACKGROUND_COLOUR = (120,120,120)
COLOURS = [pygame.Color(255, 255, 255), pygame.Color(255, 0, 0), pygame.Color(255,255,0), pygame.Color(0,255,0), pygame.Color(0,255,255), pygame.Color(0,0,255), pygame.Color(255,0,255)]
linewidth = LINE_WIDTH_DEFAULT
colour  =  pygame.Color(255, 255, 255)

COLOURINDICATOR_RECT = pygame.Rect(scr_w - 16, 0, 16, 12)

scr = pygame.display.set_mode((scr_w, scr_h))
draw_surf = pygame.Surface((scr_w - 16, scr_h - 12))
mouse_pos = (0,0)
last_mouse_pos = (0,0)

run = True
erasing = False

try:
    pygame.mouse.set_visible(0)
except pygame.error:
    pass

while run:
    for event in pygame.event.get():
        if event.type == QUIT:
            run = False
        if event.type == KEYDOWN:
            if event.key == K_ESCAPE:
                run = False
                break
            
            # my odd way of cycling through all the K_<number> pygame attributes,
            # probably really bad practice but it's kinda interesting.
            for i in range(len(COLOURS)):
                attribute_name = 'K_' + (str(i)[0])
                attribute = getattr(pygame, attribute_name)
                if event.key == attribute:
                    colour = COLOURS[i]
                    linewidth = LINE_WIDTH_DEFAULT
                    erasing = False
                    break
            # eraser (black)
            if event.key in {K_c, K_e}:
                colour = (0,0,0)
                linewidth = LINE_WIDTH_ERASER
                erasing = True
            elif event.key == K_s:
                # save image
                if os.path.exists(os.path.expanduser('~/Desktop')) \
                    and os.path.isdir(os.path.expanduser('~/Desktop')):
                    # save to desktop
                    fname = time.strftime("~/Desktop/Painting at %H-%M-%S on %d-%m-%Y.png",
                                          time.localtime(time.time()))
                else:
                    # save to $HOME
                    fname = time.strftime("~/painting_%d-%m-%Y_%H-%M-%S.png",
                                          time.localtime(time.time()))
                fname = os.path.expanduser(fname)
                pygame.image.save(draw_surf, fname)
                print("Saved to", fname)
    
    if any(pygame.mouse.get_just_pressed()):
        mouse_pos = pygame.mouse.get_pos()
        last_mouse_pos = mouse_pos
    if any(pygame.mouse.get_pressed(5)):
        pygame.draw.aacircle(draw_surf, colour, mouse_pos, linewidth / 2)
        if pygame.mouse.get_pressed()[0]:
            pygame.draw.aaline(draw_surf, colour, last_mouse_pos, mouse_pos, linewidth)
        last_mouse_pos = mouse_pos
        mouse_pos = pygame.mouse.get_pos()

    scr.fill(BACKGROUND_COLOUR)
    scr.blit(draw_surf, (0,0)) # image to draw on
    pygame.draw.rect(scr, colour, COLOURINDICATOR_RECT) # the little colour indicator in the top-right
    pygame.draw.aacircle(scr, colour if not erasing else (24,24,24),
                         pygame.mouse.get_pos(), linewidth) # apple-pencil-hover emulation :)
    pygame.display.flip()

try:
    pygame.mouse.set_visible(1)
except pygame.error:
    pass

pygame.quit() 

TODO: save function that quickly saves draw_surf to ~/Desktop/painting_{iso8601_date}_{hr}-{min}-{sec}.png when S is pressed. [DONE]

Among other things...