tetris.py

#!/usr/bin/python3
from time import sleep
from random import choice
import pygame
from pygame.locals import *

I, J, L, O, S, T, Z, N= ((0, 255, 255), (0, 0, 255), (255, 165, 0),
	(255, 255, 0), (0, 255, 0), (170, 0, 255), (255, 0, 0), None)

class tet_shape:
	lose = False
	
	def __init__(self, shape, grid=None, x=3, y=-2):
		self.shape = shape
		self.grid = grid
		self.x = x
		self.y = y
	
	def shift(self, amount):
		if amount > 0:
			self.x += 1
		elif amount < 0:
			self.x -= 1
		self.x = max(self.x, 0)
		self.x = min(self.x, len(self.grid.shape[0]) - len(self.shape[0]))
		for y, row in enumerate(self.shape):
			for x, square in enumerate(row):
				ay, ax = y + self.y, x + self.x
				if square and self.grid.shape[ay][ax]:
					self.x -= amount
					return
		if amount < 0:
			self.shift(amount + 1)
		elif amount > 0:
			self.shift(amount - 1)
	
	def rotate(self):
		self.shape = list(zip(*self.shape))[::-1]
		if self.x != max(self.x, 0) or \
			self.x != min(self.x, len(self.grid.shape[0]) - len(self.shape[0])):
				self.shape = list(zip(*self.shape[::-1]))
		for y, row in enumerate(self.shape):
			for x, square in enumerate(row):
				ay, ax = y + self.y, x + self.x
				if square and self.grid.shape[ay][ax]:
					self.shape = list(zip(*self.shape[::-1]))
	
	def drop(self):
		y = len(self.shape) + self.y
		if y <= 1:
			self.y += 1
			return None
		elif y >= len(self.grid.shape):
			self.merge()
			return True
		else:
			for y, row in enumerate(self.shape):
				for x, square in enumerate(row):
					ay, ax = y + self.y + 1, x + self.x
					if square and self.grid.shape[ay][ax]:
						self.merge()
						return True
		self.y += 1
	
	def merge(self):
		for y, row in enumerate(self.shape):
			for x, square in enumerate(row):
				ay, ax = y + self.y, x + self.x
				if square and not self.grid.shape[ay][ax]:
					self.grid.shape[ay][ax] = square
				elif square and self.grid.shape[ay][ax]:
					self.grid.lose = True
	
	def get_squares(self):
		for y, row in enumerate(self.shape):
			for x, square in enumerate(row):
				if square:
					rect = self.SQ.get_rect()
					px = (self.x * 25) + (x * 25)
					py = (self.y * 25) + (y * 25)
					rect = rect.move(px, py)
					yield (square, rect)
	
	def draw(self, surf):
		for square in self.get_squares():
			self.SQ.fill(square[0])
			surf.blit(self.SQ, square[1])
	
	SQ = pygame.Surface((25, 25))
	SHAPE = [[[I, I, I, I]],
		[[J, J, J],
		 [N, N, J]],
		[[L, L, L],
		 [L, N, N]],
		[[O, O],
		 [O, O]],
		[[N, S, S],
		 [S, S, N]],
		[[T, T, T],
		 [N, T, N]],
		[[Z, Z, N],
		 [N, Z, Z]]]

def draw(scr, shapes):
	scr.fill((0, 0, 0))
	for shape in shapes:
		shape.draw(scr)
	pygame.display.flip()

def on_xout():
	exit()

def on_keydown(event, shape):
	if event.key == K_UP:
		shape.rotate()
	elif event.key == K_DOWN:
		dropped = shape.drop()
		if dropped:
			return tet_shape(choice(tet_shape.SHAPE), grid=shape.grid)
	elif event.key == K_SPACE:
		dropped = shape.drop()
		while not dropped:
			dropped = shape.drop()
		return tet_shape(choice(tet_shape.SHAPE), grid=shape.grid)
	elif event.key == K_LEFT:
		shape.shift(-1)
	elif event.key == K_RIGHT:
		shape.shift(1)
	return shape

def on_tick(shape, grid):
	dropped = shape.drop()
	for y, row in enumerate(grid.shape):
		if False not in row:
			grid.shape.pop(y)
			grid.shape.insert(0, [False for i in range(10)])
	if dropped:
		return tet_shape(choice(tet_shape.SHAPE), grid=shape.grid)
	else:
		return shape


def main():
	pygame.init()
	
	false_grid = [[False for i in range(10)] for i in range(20)]
	grid = tet_shape(false_grid, x=0, y=0)
	shape = tet_shape(choice(tet_shape.SHAPE), grid=grid)
	
	pygame.key.set_repeat(400, 100)
	pygame.time.set_timer(USEREVENT, 500)
	scr = pygame.display.set_mode((250, 500))
	pygame.display.set_caption("Tetris")
	try:
		pygame.mixer.music.load("tetris.mp3")
		pygame.mixer.music.play(loops=-1)
	except:
		pass
	
	while not grid.lose:
		for event in pygame.event.get():
			if event.type == QUIT:
				on_xout()
			elif event.type == KEYDOWN:
				shape = on_keydown(event, shape)
			elif event.type == USEREVENT:
				shape = on_tick(shape, grid)
		draw(scr, (grid, shape))
	scr.fill((0, 0, 0))
	pygame.display.flip()
	pygame.mixer.music.fadeout(2000)
	sleep(2)
	exit()

if __name__ == "__main__":
	main()