Add piece generation and stick

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

1 files changed, 32 insertions, 13 deletions

diff --git a/cameltris.py b/cameltris.py
index a5fd6e5..b0c78e1 100755
--- a/cameltris.py
+++ b/cameltris.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import contextlib
+import random
 import sys
 import time
 
@@ -157,6 +158,31 @@ def rotate_piece_clockwise():
         current_piece.rotate_counter_clockwise()
         raise WouldCollide()
 
+
+def generate_piece():
+    piece = random.choice((TPiece, SPiece, IPiece, ZPiece, SquarePiece))()
+
+    for row_id, row in enumerate(piece.elements):
+        if list(filter(lambda x: x is not None, row)):
+            break
+
+    initial_y_position = -row_id
+    initial_x_position = (len(grid[0]) // 2) - (len(piece.elements[0]) // 2)
+
+    return (piece, [initial_y_position, initial_x_position])
+
+
+def stick_piece():
+    global current_piece, current_piece_position
+
+    for row_id, row in enumerate(current_piece.elements):
+        for col_id, element in enumerate(row):
+            if element is None:
+                continue
+            grid[row_id + current_piece_position[0]][col_id + current_piece_position[1]] = element
+
+    current_piece, current_piece_position = generate_piece()
+
 pygame.init()
 
 black = (0, 0, 0)
@@ -185,17 +211,7 @@ screen = pygame.display.set_mode((500, 1000))
 
 grid = [[None for _ in range(10)] for _ in range(20)]
 
-# Let's start with a T piece
-current_piece = TPiece()
-
-for row_id, row in enumerate(current_piece.elements):
-    if list(filter(lambda x: x is not None, row)):
-        break
-
-initial_y_position = -row_id
-initial_x_position = (len(grid[0]) // 2) - (len(current_piece.elements[0]) // 2)
-
-current_piece_position = [initial_y_position, initial_x_position]
+current_piece, current_piece_position = generate_piece()
 
 while True:
     for event in pygame.event.get():
@@ -205,8 +221,6 @@ while True:
             if event.key in (pygame.K_ESCAPE, pygame.K_q):
                 sys.exit()
             with contextlib.suppress(WouldCollide):
-                if event.key in (pygame.K_DOWN, pygame.K_j):
-                    move_piece_down()
                 if event.key in (pygame.K_UP, pygame.K_k):
                     move_piece_up()
                 if event.key in (pygame.K_RIGHT, pygame.K_l):
@@ -217,6 +231,11 @@ while True:
                     rotate_piece_clockwise()
                 if event.key == pygame.K_d:
                     rotate_piece_counter_clockwise()
+            try:
+                if event.key in (pygame.K_DOWN, pygame.K_j):
+                    move_piece_down()
+            except WouldCollide:
+                stick_piece()
 
     refresh_screen()
     pygame.display.flip()