Add snake control.

snake/game_object.py

@@ -1,6 +1,9 @@
+import pygame
+
+
 class GameObject:
-    def draw(self, surface):
+    def draw(self, surface: pygame.Surface) -> None:
         pass
 
-    def update(self):
+    def update(self, user_input: int) -> None:
         pass

snake/input_manager.py

@@ -3,15 +3,28 @@ import pygame
 
 class InputManager:
     QUIT = 0
+    IDLE = 1
+    LEFT = 2
+    RIGHT = 3
+    UP = 4
+    DOWN = 5
+
+    def __init__(self):
+        self.__key_map_dict = {
+            pygame.K_ESCAPE: InputManager.QUIT,
+            pygame.K_LEFT: InputManager.LEFT,
+            pygame.K_RIGHT: InputManager.RIGHT,
+            pygame.K_UP: InputManager.UP,
+            pygame.K_DOWN: InputManager.DOWN
+        }
 
     def process_input(self) -> int:
-        last_input = None
+        last_input = InputManager.IDLE
 
         for event in pygame.event.get():
             if event.type == pygame.QUIT:
                 last_input = InputManager.QUIT
             elif event.type == pygame.KEYDOWN:
-                if event.key == pygame.K_ESCAPE:
-                    last_input = InputManager.QUIT
+                last_input = self.__key_map_dict.get(event.key, InputManager.IDLE)
         
         return last_input

snake/main.py

@@ -11,9 +11,9 @@ def draw_all(window: Window, game_objs: list[GameObject]) -> None:
         window.draw_object(obj)
 
 
-def update_all(game_objs: list[GameObject]) -> None:
+def update_all(game_objs: list[GameObject], user_input: int) -> None:
     for obj in game_objs:
-        obj.update()
+        obj.update(user_input)
 
 
 if __name__ == '__main__':
@@ -33,7 +33,7 @@ if __name__ == '__main__':
         last_input = input_manager.process_input()
 
         draw_all(window, game_objs)
-        update_all(game_objs)
+        update_all(game_objs, last_input)
 
         clock.tick(framerate)
         pygame.display.flip()

snake/snake.py

@@ -1,6 +1,7 @@
 import pygame
 from game_object import GameObject
 from snake_segment import SnakeSegment
+from snake_state_machine import SnakeStateMachine
 
 
 class Snake(GameObject):
@@ -9,19 +10,16 @@ class Snake(GameObject):
                                         radius=radius,
                                         color=color) 
                            for idx in range(length)]
+        
+        self.__state_machine = SnakeStateMachine()
 
     def draw(self, surface: pygame.Surface) -> None:
         for segment in self.__segments:
             segment.draw(surface)
 
-    def update(self) -> None:
-        head_pos = self.__segments[0].get_position()
-        head_radius = self.__segments[0].get_radius()
-        head_color = self.__segments[0].get_color()
-
-        new_head = SnakeSegment(center=(head_pos[0]-2*head_radius, head_pos[1]),
-                                radius=head_radius,
-                                color=head_color)
+    def update(self, user_input: int=None) -> None:
+        self.__state_machine.update(user_input)
+        new_head = self.__state_machine.get_state().get_next_head(self.__segments[0])
         
         self.__segments.insert(0, new_head)
         self.__segments.pop()

snake/snake_down_movement_state.py

@@ -0,0 +1,6 @@
+from snake_state import SnakeState
+
+
+class SnakeDownMovementState(SnakeState):
+    def _get_next_pos(self, old_pos: tuple, radius: float) -> tuple:
+        return (old_pos[0], old_pos[1]+2*radius)

snake/snake_left_movement_state.py

@@ -0,0 +1,6 @@
+from snake_state import SnakeState
+
+
+class SnakeLeftMovementState(SnakeState):
+    def _get_next_pos(self, old_pos: tuple, radius: float) -> tuple:
+        return (old_pos[0]-2*radius, old_pos[1])

snake/snake_right_movement_state.py

@@ -0,0 +1,6 @@
+from snake_state import SnakeState
+
+
+class SnakeRightMovementState(SnakeState):
+    def _get_next_pos(self, old_pos: tuple, radius: float) -> tuple:
+        return (old_pos[0]+2*radius, old_pos[1])

snake/snake_state.py

@@ -0,0 +1,11 @@
+from snake_segment import SnakeSegment
+
+
+class SnakeState:
+    def _get_next_pos(self, old_pos: tuple, radius:float) -> tuple:
+        pass
+
+    def get_next_head(self, old_head: SnakeSegment) -> SnakeSegment:
+        return SnakeSegment(center=self._get_next_pos(old_head.get_position(), old_head.get_radius()),
+                            radius=old_head.get_radius(),
+                            color=old_head.get_color())

snake/snake_state_machine.py

@@ -0,0 +1,25 @@
+from snake_up_movement_state import SnakeUpMovementState
+from snake_down_movement_state import SnakeDownMovementState
+from snake_left_movement_state import SnakeLeftMovementState
+from snake_right_movement_state import SnakeRightMovementState
+from input_manager import InputManager
+
+
+class SnakeStateMachine:
+    def __init__(self):
+        self.__state = SnakeLeftMovementState()
+
+        self.__next_state_dict = {
+            InputManager.UP: SnakeUpMovementState(),
+            InputManager.DOWN: SnakeDownMovementState(),
+            InputManager.LEFT: SnakeLeftMovementState(),
+            InputManager.RIGHT: SnakeRightMovementState()
+        }
+
+    
+    def update(self, user_input):
+        self.__state = self.__next_state_dict.get(user_input, self.__state)
+
+
+    def get_state(self):
+        return self.__state

snake/snake_up_movement_state.py

@@ -0,0 +1,6 @@
+from snake_state import SnakeState
+
+
+class SnakeUpMovementState(SnakeState):
+    def _get_next_pos(self, old_pos: tuple, radius: float) -> tuple:
+        return (old_pos[0], old_pos[1]-2*radius)