Initial commit

snake.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, replace
+
+from directions import Direction
+
+
+@dataclass
+class Block:
+    x: int
+    y: int
+
+    def move(self, direction: Direction) -> None:
+        """
+        Moves the block in the given direction.
+        """
+        if direction is Direction.NORTH:
+            self.y -= 1
+        elif direction is Direction.EAST:
+            self.x += 1
+        elif direction is Direction.SOUTH:
+            self.y += 1
+        elif direction is Direction.WEST:
+            self.x -= 1
+
+    def to_tuple(self) -> tuple:
+        """
+        Returns the block in tuple representation (x,y).
+        """
+        return (self.x, self.y)
+
+    def __repr__(self) -> str:
+        return f"Block({self.x}, {self.y})"
+
+
+class Snake:
+    def __init__(self, x: int, y: int):
+        self.body = [Block(x, y)]
+
+    def preview(self, direction: Direction) -> Block:
+        """
+        Previews the head of the snake after a move in the given direction is made.
+        """
+        head = replace(self.body[0])
+        head.move(direction)
+        return head
+
+    def move(self, direction: Direction, extend: bool = False) -> None:
+        """
+        Moves the snake in the given direction. Extends the snake by one block if extend
+        is true.
+        """
+        head = replace(self.body[0])
+        head.move(direction)
+        self.body.insert(0, head)
+
+        if not extend:
+            self.body.pop()
+
+    def __len__(self) -> int:
+        return len(self.body)
+
+    def __contains__(self, item: Block) -> bool:
+        return item in self.body