day15: tests

day15/day15.py

@@ -1,10 +1,14 @@
-from dataclasses import dataclass, field
-from enum import Enum
+from day15.lib.classes import AddRemove, Box, Lens, Step
 
+INPUT = "day15/input.txt"
+INPUT_SMALL = "day15/input-small.txt"
 
-def get_input() -> str:
-    with open("day15/input.txt") as file:
-        return file.read()
+
+def get_input(path: str) -> list[str]:
+    with open(path) as file:
+        data = file.read()
+    raw_steps = data.split(",")
+    return raw_steps
 
 
 def get_string_hash(string: str) -> int:
@@ -18,95 +22,23 @@ def get_string_hash(string: str) -> int:
     return value
 
 
-class AddRemove(Enum):
-    """Simple instruction to add or remove lens"""
-
-    Add = 0
-    Remove = 1
-
-
-@dataclass
-class Step:
-    """well defined step"""
-
-    lens_name: str
-    box: int
-    focal_length: int | None = None
-    process: AddRemove | None = None
-
-
-@dataclass
-class Lens:
-    """Lens object"""
-
-    name: str
-    focal_length: int
-
-    def __hash__(self) -> int:
-        return hash(str(self.name) + ":" + str(self.focal_length))
-
-    def __str__(self) -> str:
-        return f"[{self.name} {self.focal_length}]"
-
-
-@dataclass
-class Box:
-    id: int = 0
-    contents: list[Lens] = field(default_factory=list)
-
-    def add_lens(self, lens: Lens) -> None:
-        """
-        If a lens name already exists, swap its power;
-        otherwise just add it
-        """
-        for existing_lens in self.contents:
-            if lens.name == existing_lens.name:
-                existing_lens.focal_length = lens.focal_length
-                return
-        self.contents.append(lens)
-
-    def remove_lens(self, lens_name: str) -> None:
-        """if a lens with a matching name is inside, remove it"""
-        to_remove = None
-        for existing_lens in self.contents:
-            if existing_lens.name == lens_name:
-                to_remove = existing_lens
-                break
-        if to_remove is not None:
-            self.contents.remove(to_remove)
-
-    def __str__(self) -> str:
-        return f"Box {self.id}: " + " ".join(str(lens) for lens in self.contents)
-
-    def calculate_power(self) -> int:
-        """Calculates power of the box by summing powers of the lenses"""
-        result = 0
-        for slot_number, lens in enumerate(self.contents):
-            box_power = 1 + self.id
-            slot_power = slot_number + 1
-            power = box_power * slot_power * lens.focal_length
-            result += power
-
-        return result
-
-
 def parse_step_pt2(raw_step: str) -> Step:
     """Handles as step in part 2"""
     if len(splits := raw_step.split("=")) == 2:
         box = get_string_hash(splits[0])
         strength = int(splits[1].strip())
-        return Step(splits[0], box, strength, AddRemove.Add)
+        return Step(splits[0], box, AddRemove.Add, strength)
     elif len(splits := raw_step.split("-")) == 2:
         box = get_string_hash(splits[0])
-        return Step(splits[0], box, process=AddRemove.Remove)
+        return Step(splits[0], box, AddRemove.Remove)
 
     raise ValueError(raw_step)
 
 
 def process_steps_pt2(steps: list[Step]) -> int:
     """Process a list of steps"""
     boxes: list[Box] = [Box(i) for i in range(256)]
-
+    print(boxes[0])
     for step in steps:
         if step.process == AddRemove.Remove:
             boxes[step.box].remove_lens(step.lens_name)
@@ -119,16 +51,24 @@ def process_steps_pt2(steps: list[Step]) -> int:
     return sum(box.calculate_power() for box in boxes)
 
 
+def question1(raw_steps: list[str]) -> int:
+    return sum(get_string_hash(raw_step) for raw_step in raw_steps)
+
+
+def question2(raw_steps: list[str]) -> int:
+    steps = [parse_step_pt2(raw_step) for raw_step in raw_steps]
+    return process_steps_pt2(steps)
+
+
 def main() -> None:
     """main function"""
-    chars = get_input()
-    raw_steps = chars.split(",")
+    raw_steps = get_input(INPUT)
+
     # q1
-    print(sum(get_string_hash(raw_step) for raw_step in raw_steps))
+    print(question1(raw_steps))
 
     # q2
-    steps = [parse_step_pt2(raw_step) for raw_step in raw_steps]
-    print(process_steps_pt2(steps))
+    print(question2(raw_steps))
 
 
 if __name__ == "__main__":

day15/lib/classes.py

@@ -0,0 +1,75 @@
+from dataclasses import dataclass, field
+from enum import IntEnum
+
+
+class AddRemove(IntEnum):
+    """Simple instruction to add or remove lens"""
+
+    Add = 0
+    Remove = 1
+
+
+@dataclass
+class Step:
+    """well defined step"""
+
+    lens_name: str
+    box: int
+    process: AddRemove
+    focal_length: int | None = None
+
+
+@dataclass
+class Lens:
+    """Lens object"""
+
+    name: str
+    focal_length: int
+
+    def __hash__(self) -> int:
+        return hash(str(self.name) + ":" + str(self.focal_length))
+
+    def __str__(self) -> str:
+        return f"[{self.name} {self.focal_length}]"
+
+
+@dataclass
+class Box:
+    id: int = 0
+    contents: list[Lens] = field(default_factory=list)
+
+    def add_lens(self, lens: Lens) -> None:
+        """
+        If a lens name already exists, swap its power;
+        otherwise just add it
+        """
+        print(self.contents)
+        for existing_lens in self.contents:
+            if lens.name == existing_lens.name:
+                existing_lens.focal_length = lens.focal_length
+                return
+        self.contents.append(lens)
+
+    def remove_lens(self, lens_name: str) -> None:
+        """if a lens with a matching name is inside, remove it"""
+        to_remove = None
+        for existing_lens in self.contents:
+            if existing_lens.name == lens_name:
+                to_remove = existing_lens
+                break
+        if to_remove is not None:
+            self.contents.remove(to_remove)
+
+    def __str__(self) -> str:
+        return f"Box {self.id}: " + " ".join(str(lens) for lens in self.contents)
+
+    def calculate_power(self) -> int:
+        """Calculates power of the box by summing powers of the lenses"""
+        result = 0
+        for slot_number, lens in enumerate(self.contents):
+            box_power = 1 + self.id
+            slot_power = slot_number + 1
+            power = box_power * slot_power * lens.focal_length
+            result += power
+
+        return result

day15/tests/test_classes.py

@@ -0,0 +1,24 @@
+from day15.lib.classes import Box, Lens
+
+
+def test_box() -> None:
+    box = Box(0)
+    rn = Lens("rn", 1)
+    box.add_lens(rn)
+    assert box.contents == [rn]
+    cm = Lens("cm", 2)
+    box.add_lens(cm)
+    assert box.contents == [rn, cm]
+
+    box = Box(3)
+    box.add_lens(pc := Lens("pc", 4))
+    box.add_lens(ot := Lens("ot", 9))
+    box.add_lens(ab := Lens("ab", 5))
+    assert box.contents == [pc, ot, ab]
+    box.remove_lens("pc")
+    assert box.contents == [ot, ab]
+    box.add_lens(pc := Lens("pc", 6))
+    assert box.contents == [ot, ab, pc]
+    box.add_lens(ot2 := Lens("ot", 7))
+    assert box.contents == [ot2, ab, pc]
+    assert box.contents[0].focal_length == 7

day15/tests/test_day15.py

@@ -0,0 +1,43 @@
+from day15.day15 import (
+    INPUT_SMALL,
+    get_input,
+    get_string_hash,
+    parse_step_pt2,
+    question1,
+    question2,
+)
+from day15.lib.classes import AddRemove, Step
+
+
+def test_get_input() -> None:
+    steps: list[str] = get_input(INPUT_SMALL)
+    assert len(steps) == 11
+    assert steps[0] == "rn=1"
+
+
+def test_parse_pt2() -> None:
+    steps: list[str] = get_input(INPUT_SMALL)
+    step: Step = parse_step_pt2(steps[0])
+    assert step.lens_name == "rn"
+    assert step.process == AddRemove.Add
+
+
+def test_questions() -> None:
+    steps: list[str] = get_input(INPUT_SMALL)
+    assert question1(steps) == 1320
+    assert question2(steps) == 145
+
+
+def test_get_string_hash() -> None:
+    assert get_string_hash("rn=1") == 30
+    assert get_string_hash("rn=1") == 30
+    assert get_string_hash("cm-") == 253
+    assert get_string_hash("qp=3") == 97
+    assert get_string_hash("cm=2") == 47
+    assert get_string_hash("qp-") == 14
+    assert get_string_hash("pc=4") == 180
+    assert get_string_hash("ot=9") == 9
+    assert get_string_hash("ab=5") == 197
+    assert get_string_hash("pc-") == 48
+    assert get_string_hash("pc=6") == 214
+    assert get_string_hash("ot=7") == 231