63 lines
1.5 KiB
Python
63 lines
1.5 KiB
Python
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
|