Fix 2022 day 16 for progress API.

src/holt59/aoc/2022/day16.py

@@ -6,8 +6,6 @@ import re
 from collections import defaultdict
 from typing import Any, FrozenSet, Iterator, NamedTuple
 
-from tqdm import tqdm
-
 from ..base import BaseSolver
 
 
@@ -62,68 +60,70 @@ def update_with_better(
     node_at_times[flowing] = max(node_at_times[flowing], flow)
 
 
-def part_1(
-    start_pipe: Pipe,
-    max_time: int,
-    distances: dict[tuple[Pipe, Pipe], int],
-    relevant_pipes: FrozenSet[Pipe],
-):
-    node_at_times: dict[int, dict[Pipe, dict[FrozenSet[Pipe], int]]] = defaultdict(
-        lambda: defaultdict(lambda: defaultdict(lambda: 0))
-    )
-    node_at_times[0] = {start_pipe: {frozenset(): 0}}
-
-    for time in range(max_time):
-        for c_pipe, nodes in node_at_times[time].items():
-            for flowing, flow in nodes.items():
-                for target in relevant_pipes:
-                    distance = distances[c_pipe, target] + 1
-                    if time + distance >= max_time or target in flowing:
-                        continue
-
-                    update_with_better(
-                        node_at_times[time + distance][target],
-                        flow + sum(pipe.flow for pipe in flowing) * distance,
-                        flowing | {target},
-                    )
-
-                update_with_better(
-                    node_at_times[max_time][c_pipe],
-                    flow + sum(pipe.flow for pipe in flowing) * (max_time - time),
-                    flowing,
-                )
-
-    return max(
-        flow
-        for nodes_of_pipe in node_at_times[max_time].values()
-        for flow in nodes_of_pipe.values()
-    )
-
-
-def part_2(
-    start_pipe: Pipe,
-    max_time: int,
-    distances: dict[tuple[Pipe, Pipe], int],
-    relevant_pipes: FrozenSet[Pipe],
-):
-    def compute(pipes_for_me: FrozenSet[Pipe]) -> int:
-        return part_1(start_pipe, max_time, distances, pipes_for_me) + part_1(
-            start_pipe, max_time, distances, relevant_pipes - pipes_for_me
-        )
-
-    combs = [
-        frozenset(relevant_pipes_1)
-        for r in range(2, len(relevant_pipes) // 2 + 1)
-        for relevant_pipes_1 in itertools.combinations(relevant_pipes, r)
-    ]
-
-    return max(compute(comb) for comb in tqdm(combs))
-
-
 # === MAIN ===
 
 
 class Solver(BaseSolver):
+    def part_1(
+        self,
+        start_pipe: Pipe,
+        max_time: int,
+        distances: dict[tuple[Pipe, Pipe], int],
+        relevant_pipes: FrozenSet[Pipe],
+    ):
+        node_at_times: dict[int, dict[Pipe, dict[FrozenSet[Pipe], int]]] = defaultdict(
+            lambda: defaultdict(lambda: defaultdict(lambda: 0))
+        )
+        node_at_times[0] = {start_pipe: {frozenset(): 0}}
+
+        for time in range(max_time):
+            for c_pipe, nodes in node_at_times[time].items():
+                for flowing, flow in nodes.items():
+                    for target in relevant_pipes:
+                        distance = distances[c_pipe, target] + 1
+                        if time + distance >= max_time or target in flowing:
+                            continue
+
+                        update_with_better(
+                            node_at_times[time + distance][target],
+                            flow + sum(pipe.flow for pipe in flowing) * distance,
+                            flowing | {target},
+                        )
+
+                    update_with_better(
+                        node_at_times[max_time][c_pipe],
+                        flow + sum(pipe.flow for pipe in flowing) * (max_time - time),
+                        flowing,
+                    )
+
+        return max(
+            flow
+            for nodes_of_pipe in node_at_times[max_time].values()
+            for flow in nodes_of_pipe.values()
+        )
+
+    def part_2(
+        self,
+        start_pipe: Pipe,
+        max_time: int,
+        distances: dict[tuple[Pipe, Pipe], int],
+        relevant_pipes: FrozenSet[Pipe],
+    ):
+        def compute(pipes_for_me: FrozenSet[Pipe]) -> int:
+            return self.part_1(
+                start_pipe, max_time, distances, pipes_for_me
+            ) + self.part_1(
+                start_pipe, max_time, distances, relevant_pipes - pipes_for_me
+            )
+
+        combs = [
+            frozenset(relevant_pipes_1)
+            for r in range(2, len(relevant_pipes) // 2 + 1)
+            for relevant_pipes_1 in itertools.combinations(relevant_pipes, r)
+        ]
+
+        return max(compute(comb) for comb in self.progress.wrap(combs))
+
     def solve(self, input: str) -> Iterator[Any]:
         lines = [line.strip() for line in input.splitlines()]
 
@@ -153,7 +153,7 @@ class Solver(BaseSolver):
         relevant_pipes = frozenset(pipe for pipe in pipes.values() if pipe.flow > 0)
 
         # 1651, 1653
-        yield part_1(pipes["AA"], 30, distances, relevant_pipes)
+        yield self.part_1(pipes["AA"], 30, distances, relevant_pipes)
 
         # 1707, 2223
-        yield part_2(pipes["AA"], 26, distances, relevant_pipes)
+        yield self.part_2(pipes["AA"], 26, distances, relevant_pipes)