improved day05

advent/days/day05/solution.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass
 
-from typing import Iterator
+from typing import ClassVar, Iterator
 
 from advent.parser.parser import P
 
@@ -9,97 +9,91 @@ day_num = 5
 
 
 def part1(lines: Iterator[str]) -> str:
-    state = State.parse(lines)
+    crane = Crane.parse(lines, False)
-    crates = state.all_moves9000()
+    crates = crane.perform_all_moves()
-    return State.top(crates)
+    return Crane.top(crates)
 
 
 def part2(lines: Iterator[str]) -> str:
-    state = State.parse(lines)
+    crane = Crane.parse(lines, True)
-    crates = state.all_moves9001()
+    crates = crane.perform_all_moves()
-    return State.top(crates)
+    return Crane.top(crates)
 
 
-crate_parser: P[str | None] = P.either(P.any_char().in_brackets(), P.string("   ").replace(None))
+@dataclass(slots=True, frozen=True)
-crate_row_parser = crate_parser.sep_by(P.is_char(' '))
+class Move:
-amount_parser = P.snd(P.string("move "), P.unsigned())
+    amount: int
-from_parser = P.snd(P.string(" from "), P.unsigned())
+    frm: int
-to_parser = P.snd(P.string(" to "), P.unsigned())
+    to: int
-move_parser = P.seq(amount_parser, from_parser, to_parser)
 
-
+    amount_parser: ClassVar[P[int]] = P.snd(P.string("move "), P.unsigned())
-@dataclass(slots=True)
+    from_parser: ClassVar[P[int]] = P.snd(P.string(" from "), P.unsigned())
-class State:
+    to_parser: ClassVar[P[int]] = P.snd(P.string(" to "), P.unsigned())
-    crates: list[str]
+    move_parser: ClassVar[P[tuple[int, int, int]]] = P.seq(amount_parser, from_parser, to_parser)
-    moves: list[tuple[int, int, int]]
 
     @staticmethod
-    def parse_crate_line(line: str) -> list[None | str] | None:
+    def parse(line: str) -> Move:
-        return crate_row_parser.parse(line).get_or(None)
+        amount, frm, to = Move.move_parser.parse(line).get()
+        return Move(amount, frm - 1, to - 1)
+
+    def do_move(self, crates: list[str], as_9001: bool) -> list[str]:
+        """
+        Moves the given crates by the provided move. Will fail if there are not enough crates
+        in the from stack
+        """
+        if as_9001:
+            crates[self.to] += crates[self.frm][-self.amount:]
+        else:
+            crates[self.to] += crates[self.frm][-self.amount:][::-1]
+        crates[self.frm] = crates[self.frm][:-self.amount]
+        return crates
+
+
+@dataclass(slots=True, frozen=True)
+class Crane:
+    stacks: list[str]
+    moves: list[Move]
+    is_9001: bool
+
+    crate_parser: ClassVar[P[str | None]] = P.either(
+        P.any_char().in_brackets(), P.string("   ").replace(None))
+    crate_row_parser: ClassVar[P[list[str | None]]] = crate_parser.sep_by(P.is_char(' '))
+
+    @staticmethod
+    def parse_crate_row(line: str) -> list[None | str] | None:
+        return Crane.crate_row_parser.parse(line).get_or(None)
 
     @staticmethod
     def parse_drawing(lines: Iterator[str]) -> list[str]:
-        result: list[str] = []
+        stacks: list[str] = []
         for line in lines:
-            crates = State.parse_crate_line(line)
+            crate_row = Crane.parse_crate_row(line)
-            if crates is None:
+            if crate_row is None:
-                return result
+                next(lines)  # Empty line
+                return stacks
 
-            if len(result) < len(crates):
+            if len(stacks) < len(crate_row):
-                result += [""] * (len(crates) - len(result))
+                stacks += [""] * (len(crate_row) - len(stacks))
-            for stack, crate in enumerate(crates):
+
+            for stack_num, crate in enumerate(crate_row):
                 if crate is not None:
-                    result[stack] = crate + result[stack]
+                    stacks[stack_num] = crate + stacks[stack_num]
 
         raise Exception("Can never happen")
 
     @staticmethod
-    def parse_move(line: str) -> tuple[int, int, int]:
+    def parse(lines: Iterator[str], is_9001: bool) -> Crane:
-        amount, frm, to = move_parser.parse(line).get()
+        drawing = Crane.parse_drawing(lines)
-        return amount, frm - 1, to - 1
+        moves = [Move.parse(line) for line in lines]
-
+        return Crane(drawing, moves, is_9001)
-    @staticmethod
-    def parse(lines: Iterator[str]) -> State:
-        drawing = State.parse_drawing(lines)
-        next(lines)
-        moves = [State.parse_move(line) for line in lines]
-        return State(drawing, moves)
-
-    @staticmethod
-    def do_move9000(crates: list[str], move: tuple[int, int, int]) -> list[str]:
-        """
-        Moves the given crates by the provided move. Will fail if there are not enough crates
-        in the from stack
-        """
-        for _ in range(move[0]):
-            crates[move[2]] += crates[move[1]][-1]
-            crates[move[1]] = crates[move[1]][:-1]
-
-        return crates
-
-    def all_moves9000(self) -> list[str]:
-        crates = self.crates
-        for move in self.moves:
-            crates = State.do_move9000(crates, move)
-        return crates
-
-    @staticmethod
-    def do_move9001(crates: list[str], move: tuple[int, int, int]) -> list[str]:
-        """
-        Moves the given crates by the provided move. Will fail if there are not enough crates
-        in the from stack
-        """
-        crates[move[2]] += crates[move[1]][-move[0]:]
-        crates[move[1]] = crates[move[1]][:-move[0]]
-        return crates
-
-    def all_moves9001(self) -> list[str]:
-        crates = self.crates
-        for move in self.moves:
-            crates = State.do_move9001(crates, move)
-        return crates
 
     @staticmethod
     def top(crates: list[str]) -> str:
         """ Lists the last item in the given stacks. Fails if any stack is empty """
         return ''.join(stack[-1] for stack in crates)
+
+    def perform_all_moves(self) -> list[str]:
+        stacks = self.stacks
+        for move in self.moves:
+            stacks = move.do_move(stacks, self.is_9001)
+        return stacks

advent/days/day05/test_solution.py

@@ -1,6 +1,6 @@
 from advent.common import utils
 
-from .solution import day_num, part1, part2, State
+from .solution import Move, day_num, part1, part2, Crane
 
 
 def test_part1():
@@ -20,57 +20,59 @@ def test_part2():
 def test_parse_line():
     input = "    [D]    "
     expected = [None, 'D', None]
-    result = State.parse_crate_line(input)
+    result = Crane.parse_crate_row(input)
     assert result == expected
 
 
 def test_parse_line2():
     input = "[Z] [M] [P]"
     expected = ["Z", 'M', "P"]
-    result = State.parse_crate_line(input)
+    result = Crane.parse_crate_row(input)
     assert result == expected
 
 
 def test_drawing():
     data = utils.read_data(day_num, 'test01.txt')
     expected = ["ZN", "MCD", "P"]
-    result = State.parse_drawing(data)
+    result = Crane.parse_drawing(data)
     assert result == expected
 
 
 def test_parse_move():
     input = "move 1 from 2 to 1"
-    expected = 1, 1, 0
+    expected = Move(1, 1, 0)
-    result = State.parse_move(input)
+    result = Move.parse(input)
     assert result == expected
 
 
 def test_parse_all():
     data = utils.read_data(day_num, 'test01.txt')
-    expected = State(["ZN", "MCD", "P"], [(1, 1, 0), (3, 0, 2), (2, 1, 0), (1, 0, 1)])
+    expected = Crane(
-    result = State.parse(data)
+        ["ZN", "MCD", "P"],
+        [Move(1, 1, 0), Move(3, 0, 2), Move(2, 1, 0), Move(1, 0, 1)], True)
+    result = Crane.parse(data, True)
     assert result == expected
 
 
 def test_step():
     data = utils.read_data(day_num, 'test01.txt')
-    state = State.parse(data)
+    state = Crane.parse(data, True)
     expected = ["ZND", "MC", "P"]
-    result = State.do_move9000(state.crates, state.moves[0])
+    result = state.moves[0].do_move(state.stacks, False)
     assert result == expected
 
 
 def test_all_moves():
     data = utils.read_data(day_num, 'test01.txt')
-    state = State.parse(data)
+    state = Crane.parse(data, False)
     expected = ["C", "M", "PDNZ"]
-    result = state.all_moves9000()
+    result = state.perform_all_moves()
     assert result == expected
 
 
 def test_all_moves9001():
     data = utils.read_data(day_num, 'test01.txt')
-    state = State.parse(data)
+    state = Crane.parse(data, True)
     expected = ["M", "C", "PZND"]
-    result = state.all_moves9001()
+    result = state.perform_all_moves()
     assert result == expected

advent/parser/test_parser.py

@@ -1,4 +1,4 @@
-from .parser import P  # type: ignore
+from advent.parser.parser import P
 import pytest