Rely on a function to set elements of a piece

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

1 files changed, 75 insertions, 14 deletions

diff --git a/cameltris/piece.py b/cameltris/piece.py
index 33adc37..d0e97f6 100644
--- a/cameltris/piece.py
+++ b/cameltris/piece.py
@@ -1,6 +1,6 @@
 """ Module that defines all types of pieces """
 
-from typing import Optional
+from typing import Optional, Iterable
 
 import pygame
 
@@ -28,13 +28,24 @@ class Piece:
         self.rotate_clockwise()
         self.rotate_clockwise()
 
+    def set_elements_from_description(self, desc: Iterable[Iterable]):
+        """ Set each element according to the description given as an iterable of iterable.
+        Each value from description that evaluates to True produces a square.
+        Each value from description that evaluates to False produces a None value.
+        """
+        self.elements = list(map(lambda row: list(map(lambda x: self.square if x else None, row)), desc))
+
 
 class ZPiece(Piece):
     def __init__(self):
         super().__init__()
 
         self.square.fill(Color.BLUE.value)
-        self.elements = (self.square, self.square, None), (None, self.square, self.square), (None, None, None)
+        self.set_elements_from_description([
+            (1, 1, 0),
+            (0, 1, 1),
+            (0, 0, 0),
+        ])
         self.vertical = False
 
     def rotate_clockwise(self):
@@ -45,9 +56,17 @@ class ZPiece(Piece):
 
     def rotate(self):
         if self.vertical:
-            self.elements = (self.square, self.square, None), (None, self.square, self.square), (None, None, None)
+            self.set_elements_from_description([
+                (1, 1, 0),
+                (0, 1, 1),
+                (0, 0, 0),
+            ])
         else:
-            self.elements = (None, None, self.square), (None, self.square, self.square), (None, self.square, None)
+            self.set_elements_from_description([
+                (0, 0, 1),
+                (0, 1, 1),
+                (0, 1, 0),
+            ])
 
         self.vertical = not self.vertical
 
@@ -57,7 +76,11 @@ class SPiece(Piece):
         super().__init__()
 
         self.square.fill(Color.GREEN.value)
-        self.elements = (None, None, None), (None, self.square, self.square), (self.square, self.square, None)
+        self.set_elements_from_description([
+            (0, 0, 0),
+            (0, 1, 1),
+            (1, 1, 0),
+        ])
         self.vertical = False
 
     def rotate_clockwise(self):
@@ -68,9 +91,17 @@ class SPiece(Piece):
 
     def rotate(self):
         if self.vertical:
-            self.elements = (None, self.square, self.square), (self.square, self.square, None), (None, None, None)
+            self.set_elements_from_description([
+                (0, 1, 1),
+                (1, 1, 0),
+                (0, 0, 0),
+            ])
         else:
-            self.elements = (None, self.square, None), (None, self.square, self.square), (None, None, self.square)
+            self.set_elements_from_description([
+                (0, 1, 0),
+                (0, 1, 1),
+                (0, 0, 1),
+            ])
 
         self.vertical = not self.vertical
 
@@ -80,7 +111,10 @@ class SquarePiece(Piece):
         super().__init__()
 
         self.square.fill(Color.BROWN.value)
-        self.elements = ((self.square, self.square), (self.square, self.square))
+        self.set_elements_from_description([
+            (1, 1),
+            (1, 1),
+        ])
 
 
 class IPiece(Piece):
@@ -88,7 +122,12 @@ class IPiece(Piece):
         super().__init__()
 
         self.square.fill(Color.RED.value)
-        self.elements = (None, None, None, None), (None, None, None, None), (self.square, self.square, self.square, self.square), (None, None, None, None)
+        self.set_elements_from_description([
+            (0, 0, 0, 0),
+            (0, 0, 0, 0),
+            (1, 1, 1, 1),
+            (0, 0, 0, 0),
+        ])
         self.vertical = False
 
     def rotate_clockwise(self):
@@ -99,9 +138,19 @@ class IPiece(Piece):
 
     def rotate(self):
         if self.vertical:
-            self.elements = (None, None, None, None), (None, None, None, None), (self.square, self.square, self.square, self.square), (None, None, None, None)
+            self.set_elements_from_description([
+                (0, 0, 0, 0),
+                (0, 0, 0, 0),
+                (1, 1, 1, 1),
+                (0, 0, 0, 0),
+            ])
         else:
-            self.elements = (None, None, self.square, None), (None, None, self.square, None), (None, None, self.square, None), (None, None, self.square, None)
+            self.set_elements_from_description([
+                (0, 0, 1, 0),
+                (0, 0, 1, 0),
+                (0, 0, 1, 0),
+                (0, 0, 1, 0),
+            ])
 
         self.vertical = not self.vertical
 
@@ -111,7 +160,11 @@ class LPiece(Piece):
         super().__init__()
 
         self.square.fill(Color.CYAN.value)
-        self.elements = (None, None, None), (self.square, self.square, self.square), (None, None, self.square)
+        self.set_elements_from_description([
+            (0, 0, 0),
+            (1, 1, 1),
+            (0, 0, 1),
+        ])
 
 
 class JPiece(Piece):
@@ -119,7 +172,11 @@ class JPiece(Piece):
         super().__init__()
 
         self.square.fill(Color.PURPLE.value)
-        self.elements = (None, None, None), (self.square, self.square, self.square), (self.square, None, None)
+        self.set_elements_from_description([
+            (0, 0, 0),
+            (1, 1, 1),
+            (1, 0, 0),
+        ])
 
 
 class TPiece(Piece):
@@ -127,4 +184,8 @@ class TPiece(Piece):
         super().__init__()
 
         self.square.fill(Color.YELLOW.value)
-        self.elements = (None, None, None), (self.square, self.square, self.square), (None, self.square, None)
+        self.set_elements_from_description([
+            (0, 0, 0),
+            (1, 1, 1),
+            (0, 1, 0),
+        ])