Advent of Code 2024 day 15 part 1

from dataclasses import dataclass
from enum import Enum
from argparse import ArgumentParser, FileType


@dataclass(frozen=True)
class Item:
    '''An item in the warehouse.'''
    x: int
    y: int
    pushable: bool

    @classmethod
    def from_char(cls, x, y, char):
        '''Construct an Item based on the character at the given position in
           the grid.'''
        assert char in 'O#'
        pushable = (char == 'O')
        return cls(x, y, pushable)

    @property
    def gps_coords(self):
        '''The gps coordinates of this Item.'''
        return self.x + 100 * self.y


class Direction(Enum):
    '''A cardinal direction.'''
    UP = (0, -1)
    LEFT = (1, 0)
    DOWN = (0, 1)
    RIGHT = (-1, 0)

    @classmethod
    def from_char(cls, char):
        '''Construct a direction from its corresponding character.'''
        if char == '^':
            return cls.UP
        elif char == '>':
            return cls.LEFT
        elif char == 'v':
            return cls.DOWN
        elif char == '<':
            return cls.RIGHT
        else:
            raise ValueError(f'invalid direction: "{char}"')

    def move(self, x, y):
        '''Translate the coordinates a distance of 1 in this direction.'''
        dx, dy = self.value
        return x + dx, y + dy

    def move_item(self, item):
        '''Translate the item a distance of 1 in this direction.'''
        dx, dy = self.value
        return Item(item.x + dx, item.y + dy, item.pushable)


def parse_room(src):
    '''Parse the room portion of the input into a set of Items and the starting
       coordinates of the robot.'''
    lines = src.split('\n')
    robot = None
    room = set()
    for y, line in enumerate(lines):
        assert line != ''

        for x, char in enumerate(line):
            if char == '.':
                continue
            elif char == '@':
                robot = (x, y)
            else:
                room.add(Item.from_char(x, y, char))

    assert robot is not None

    return robot, room


def parse(src):
    '''Parse the input into a set of Items, the starting coordinates of the
       robot, and an iterable of Directions.'''
    room, moves = src.split('\n\n', 1)
    return (*parse_room(room),
            map(Direction.from_char, moves.replace('\n', '')))


def empty(x, y, room):
    '''Return True if there is no item at the given coordinates.'''
    return Item(x, y, True) not in room and Item(x, y, False) not in room


def print_room(robot, room):
    '''Output a string representation of the room and robot to stdout.'''
    width = max(room, key=lambda item: item.x).x + 1
    height = max(room, key=lambda item: item.y).y + 1

    for y in range(height):
        for x in range(width):
            if Item(x, y, True) in room:
                print('O', end='')
            elif Item(x, y, False) in room:
                print('#', end='')
            elif (x, y) == robot:
                print('@', end='')
            else:
                print('.', end='')
        print()
    print()


def move(robot, room, direction):
    '''Move the robot in direction, potentially pushing some boxes and updating
       the room.'''
    nx, ny = direction.move(*robot)

    if empty(nx, ny, room):  # free to move
        return (nx, ny)

    if Item(nx, ny, False) in room:  # blockage, can't do anything
        return robot

    # get the stack of boxes
    boxes = set()
    while (box := Item(nx, ny, True)) in room:
        boxes.add(box)
        nx, ny = direction.move(nx, ny)

    # if the next space is free, move them
    if empty(nx, ny, room):
        robot = direction.move(*robot)
        new_boxes = {direction.move_item(box) for box in boxes}
        room -= boxes
        room |= new_boxes

    return robot


def main(robot, room, moves, verbose=False):
    '''Perform each movement and return the sum of the gps coordinates of each
       box.'''
    if verbose:
        print_room(robot, room)

    for direction in moves:
        robot = move(robot, room, direction)
        if verbose:
            print_room(robot, room)

    return sum(item.gps_coords for item in room if item.pushable)


arg_parser = ArgumentParser()
arg_parser.add_argument('src', type=FileType('r'))
arg_parser.add_argument('-v', '--verbose', action='store_true')

if __name__ == '__main__':
    args = arg_parser.parse_args()
    print(main(*parse(args.src.read()), args.verbose))