Mars Rover

mars_rover.py

"""
Mars Rover
A simple program in Python that takes in commands and moves one or more robots around Mars.
    - The world should be modelled as a grid with size m x n
    - Your program should read the input, update the robots, and print out the final states
    of the robots
    - Each robot has a position (x, y), and an orientation (N, E, S, W)
    - Each robot can move forward one space (F), rotate left by 90 degrees (L), or rotate
    right by 90 degrees (R)
    - If a robot moves off the grid, it is marked as ‘lost’ and its last valid grid position and
    orientation is recorded
    - Going from x -> x + 1 is in the easterly direction, and y -> y + 1 is in the northerly
    direction. i.e. (0, 0) represents the south-west corner of the grid
"""
from collections import namedtuple
from enum import IntEnum
from typing import NamedTuple, Tuple, Union

import unittest


Grid = namedtuple('Grid', ['x', 'y'])


class Direction(IntEnum):
    N = 0
    E = 1
    S = 2
    W = 3


class Location(NamedTuple):
    x: int
    y: int
    direction: Direction

    def __str__(self):
        return f'({self.x}, {self.y}, {self.direction.name})'
    
    # TODO: not desirable, but makes the output pretty
    def __repr__(self):
        return self.__str__()


class Robot:
    def __init__(self, x, y, direction):
        self.x = x
        self.y = y
        self.direction = Direction[direction]
    
    def L(self):
        """Rotate left by 90 degrees."""
        new_direction_value = (self.direction.value - 1) % 4
        self.direction = Direction(new_direction_value)
    
    def R(self):
        """Rotate right by 90 degrees."""
        new_direction_value = (self.direction.value + 1) % 4
        self.direction = Direction(new_direction_value)
    
    def F(self):
        """
        Move forward one space.
        
        Going from x -> x + 1 is in the easterly direction, and
        y -> y + 1 is in the northerly direction.
        """
        if self.direction == Direction.E:
            self.x += 1
        elif self.direction == Direction.W:
            self.x -= 1
        
        elif self.direction == Direction.N:
            self.y += 1
        elif self.direction == Direction.S:
            self.y -= 1

    def get_location(self):
        """Returns the current location of the robot."""
        return Location(self.x, self.y, self.direction)


def move_robot(grid: Grid, robot: Robot, instructions: str) -> Tuple[Location, Union[None, str]]:
    """
    Moves the robot based on the instructions within the grid.
    
    Returns a tuple containing the final location of the robot and the status (None if not lost, 'LOST' if lost).
    """
    for instruction in instructions:
        if instruction not in 'LRF':
            raise ValueError(f'{repr(instruction)} is an invalid instruction. Must be on of: L, R or F.')
        loc = robot.get_location()
        getattr(robot, instruction)()
        if robot.x < 0 or robot.x > grid.x or robot.y < 0 or robot.y > grid.y:
            return loc, 'LOST'
    
    return robot.get_location(), None


class TestRobotProgram(unittest.TestCase):

    def test_move_robot_within_grid(self):
        grid = Grid(4, 8)
        robot = Robot(2, 3, 'E')
        instructions = 'LFRFF'

        final_location, status = move_robot(grid, robot, instructions)

        expected_location = Location(4, 4, Direction.E)

        self.assertEqual(final_location, expected_location)
        self.assertIsNone(status)

    def test_move_robot_outside_grid(self):
        grid = Grid(4, 8)
        robot = Robot(0, 2, 'N')
        instructions = 'FFLFRFF'

        final_location, status = move_robot(grid, robot, instructions)

        expected_location = Location(0, 4, Direction.W)

        self.assertEqual(final_location, expected_location)
        self.assertEqual('LOST', status)

    def test_invalid_instructions(self):
        grid = Grid(4, 8)
        robot = Robot(2, 3, 'E')
        invalid_instructions = 'XYZ'

        with self.assertRaisesRegex(ValueError, "'X' is an invalid instruction. Must be on of: L, R or F."):
            move_robot(grid, robot, invalid_instructions)


def run_tests():
    unittest.main()


def run_examples():
    grid = Grid(4, 8)
    print(move_robot(grid, Robot(2, 3, 'N'), 'FLLFR'))
    print(move_robot(grid, Robot(1, 0, 'S'), 'FFRLF'))


if __name__ == '__main__':
    # run_tests()
    run_examples()