Add the T piece

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

1 files changed, 23 insertions, 2 deletions

diff --git a/cameltris.py b/cameltris.py
index d2ed6f8..fdb0845 100755
--- a/cameltris.py
+++ b/cameltris.py
@@ -64,6 +64,23 @@ class IPiece(Piece):
         self.vertical = not self.vertical
 
 
+class TPiece(Piece):
+    def __init__(self):
+        self.elements = [[yellow_square, yellow_square, yellow_square], [None, yellow_square, None], [None, None, None]]
+
+    def rotate_clockwise(self):
+        # Set the corners
+        self.elements[0][0], self.elements[0][2], self.elements[2][2], self.elements[2][0] = self.elements[2][0], self.elements[0][0], self.elements[0][2], self.elements[2][2]
+        # Set the middle squares
+        self.elements[0][1], self.elements[1][2], self.elements[2][1], self.elements[1][0] = self.elements[1][0], self.elements[0][1], self.elements[1][2], self.elements[2][1]
+
+    def rotate_counter_clockwise(self):
+        # Set the corners
+        self.elements[2][0], self.elements[0][0], self.elements[0][2], self.elements[2][2] = self.elements[0][0], self.elements[0][2], self.elements[2][2], self.elements[2][0]
+        # Set the middle squares
+        self.elements[1][0], self.elements[0][1], self.elements[1][2], self.elements[2][1] = self.elements[0][1], self.elements[1][2], self.elements[2][1], self.elements[1][0]
+
+
 def refresh_screen():
     screen.fill(black)
 
@@ -85,6 +102,7 @@ white = (0xff, 0xff, 0xff)
 blue = (0x10, 0x20, 0xbb)
 green = (0x20, 0xbb, 0x10)
 red = (0xbb, 0x10, 0x20)
+yellow = (0xab, 0xd0, 0x20)
 
 white_square = pygame.Surface((48, 48))
 white_square.fill(white)
@@ -98,12 +116,15 @@ green_square.fill(green)
 red_square = pygame.Surface((48, 48))
 red_square.fill(red)
 
+yellow_square = pygame.Surface((48, 48))
+yellow_square.fill(yellow)
+
 screen = pygame.display.set_mode((500, 1000))
 
 grid = [[None for _ in range(10)] for _ in range(20)]
 
-# Let's start with a I piece
-current_piece = ZPiece()
+# 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)):