Add rotation support to Polygon

geom/geom.py

@@ -1,5 +1,7 @@
 from .math import *
 from .draw import *
+import numpy as np
+
 
 class Rect:
 
@@ -52,46 +54,75 @@ class Rect:
         pixel(surface, colour, *self.get_center())
         rect(surface, colour, self.x, self.y, self.width, self.height, width)
 
+
 class Polygon:
 
-    def __init__(self, x, y, vertices, origin_x=0, origin_y=0):
+    def __init__(self, x, y, vertices, origin_x=0, origin_y=0, rotation=0):
         self.x = x
         self.y = y
         self.vertices = vertices
+        self.length = len(vertices)
         self.origin_x = origin_x
         self.origin_y = origin_y
+        self.rotation = rotation
+        self.update_transformation()
         self.update()
 
     def update(self):
         self.translated_vertices = [
-            (self.x + self.origin_x + x,
+            (self.x + x, self.y + y)
-             self.y + self.origin_y + y)
+            for x, y in self.transformed_vertices
-            for x, y in self.vertices
         ]
+        self.segments = [
+            (self.translated_vertices[i],
+             self.translated_vertices[i + 1 if i + 1 < self.length else 0])
+            for i in range(self.length)
+        ]
+        self.normals = [
+            normalize(*lnormal(*sub(*v2, *v1)))
+            for v1, v2 in self.segments
+        ]
+
+    def update_transformation(self):
+        if self.rotation != 0:
+            rads = math.radians(self.rotation)
+            rot_matrix = np.array([[math.cos(rads), -math.sin(rads)],
+                                   [math.sin(rads), math.cos(rads)]])
+
+            self.transformed_vertices = [
+                add(*rot_matrix.dot(sub(x, y, self.origin_x, self.origin_y)),
+                    self.origin_x, self.origin_y)
+                for x, y in self.vertices
+            ]
+        else:
+            self.transformed_vertices = [
+                (self.origin_x + x, self.origin_y + y)
+                for x, y in self.vertices
+            ]
 
     def set_position(self, x, y):
         self.x = x
         self.y = y
         self.update()
 
+    def set_rotation(self, rotation):
+        self.rotation = rotation
+        self.update_transformation()
+
+    def rotate(self, degrees):
+        self.set_rotation(self.rotation + degrees)
+
     def get_normals(self):
-        return [
+        return self.normals
-            normalize(*lnormal(*sub(*v2, *v1)))
-            for v1, v2 in self.get_segments()
-        ]
 
     def get_translated_vertices(self):
         return self.translated_vertices
 
+    def get_vertices(self):
+        return self.transformed_vertices
+
     def get_segments(self):
-        cnt = len(self.vertices)
+        return self.segments
-        return [
-            (
-                self.translated_vertices[i],
-                self.translated_vertices[i + 1 if i + 1 < cnt else 0],
-            )
-            for i in range(cnt)
-        ]
 
     def draw(self, surface, colour, width):
         polygon(surface, colour, self.translated_vertices, width)