Refactor code for API (#3)

Co-authored-by: Mikael CAPELLE <mikael.capelle@thalesaleniaspace.com>
Co-authored-by: Mikaël Capelle <capelle.mikael@gmail.com>
Reviewed-on: #3

src/holt59/aoc/2022/day16.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 import heapq
 import itertools
 import re
-import sys
 from collections import defaultdict
-from typing import FrozenSet, NamedTuple
+from typing import Any, FrozenSet, Iterator, NamedTuple
 
 from tqdm import tqdm
 
+from ..base import BaseSolver
+
 
 class Pipe(NamedTuple):
     name: str
@@ -36,8 +37,8 @@ def breadth_first_search(pipes: dict[str, Pipe], pipe: Pipe) -> dict[Pipe, int]:
     Runs a BFS from the given pipe and return the shortest distance (in term of hops)
     to all other pipes.
     """
-    queue = [(0, pipe_1)]
-    visited = set()
+    queue = [(0, pipe)]
+    visited: set[Pipe] = set()
     distances: dict[Pipe, int] = {}
 
     while len(distances) < len(pipes):
@@ -122,37 +123,37 @@ def part_2(
 # === MAIN ===
 
 
-lines = sys.stdin.read().splitlines()
+class Solver(BaseSolver):
+    def solve(self, input: str) -> Iterator[Any]:
+        lines = [line.strip() for line in input.splitlines()]
 
+        pipes: dict[str, Pipe] = {}
+        for line in lines:
+            r = re.match(
+                R"Valve ([A-Z]+) has flow rate=([0-9]+); tunnels? leads? to valves? (.+)",
+                line,
+            )
+            assert r
 
-pipes: dict[str, Pipe] = {}
-for line in lines:
-    r = re.match(
-        R"Valve ([A-Z]+) has flow rate=([0-9]+); tunnels? leads? to valves? (.+)",
-        line,
-    )
-    assert r
+            g = r.groups()
 
-    g = r.groups()
+            pipes[g[0]] = Pipe(g[0], int(g[1]), g[2].split(", "))
 
-    pipes[g[0]] = Pipe(g[0], int(g[1]), g[2].split(", "))
+        # compute distances from one valve to any other
+        distances: dict[tuple[Pipe, Pipe], int] = {}
+        for pipe_1 in pipes.values():
+            distances.update(
+                {
+                    (pipe_1, pipe_2): distance
+                    for pipe_2, distance in breadth_first_search(pipes, pipe_1).items()
+                }
+            )
 
-# compute distances from one valve to any other
-distances: dict[tuple[Pipe, Pipe], int] = {}
-for pipe_1 in pipes.values():
-    distances.update(
-        {
-            (pipe_1, pipe_2): distance
-            for pipe_2, distance in breadth_first_search(pipes, pipe_1).items()
-        }
-    )
+        # valves with flow
+        relevant_pipes = frozenset(pipe for pipe in pipes.values() if pipe.flow > 0)
 
-# valves with flow
-relevant_pipes = frozenset(pipe for pipe in pipes.values() if pipe.flow > 0)
+        # 1651, 1653
+        yield part_1(pipes["AA"], 30, distances, relevant_pipes)
 
-
-# 1651, 1653
-print(part_1(pipes["AA"], 30, distances, relevant_pipes))
-
-# 1707, 2223
-print(part_2(pipes["AA"], 26, distances, relevant_pipes))
+        # 1707, 2223
+        yield part_2(pipes["AA"], 26, distances, relevant_pipes)