2024 day 20.

src/holt59/aoc/tools/graphs.py

@@ -1,5 +1,5 @@
 import heapq
-from typing import Callable, Iterable, TypeVar
+from typing import Callable, Iterable, TypeVar, overload
 
 _Node = TypeVar("_Node")
 
@@ -53,11 +53,31 @@ def make_neighbors_grid_fn(
     return _fn
 
 
+@overload
+def dijkstra(
+    start: _Node,
+    target: None,
+    neighbors: Callable[[_Node], Iterable[tuple[_Node, float]]],
+) -> dict[_Node, tuple[tuple[_Node, ...], float]]: ...
+
+
+@overload
 def dijkstra(
     start: _Node,
     target: _Node,
     neighbors: Callable[[_Node], Iterable[tuple[_Node, float]]],
-) -> tuple[tuple[_Node, ...], float] | None:
+) -> tuple[tuple[_Node, ...], float] | None: ...
+
+
+def dijkstra(
+    start: _Node,
+    target: _Node | None,
+    neighbors: Callable[[_Node], Iterable[tuple[_Node, float]]],
+) -> (
+    dict[_Node, tuple[tuple[_Node, ...], float]]
+    | tuple[tuple[_Node, ...], float]
+    | None
+):
     """
     Solve shortest-path problem using simple Dijkstra algorithm from start to target,
     using the given neighbors function.
@@ -92,4 +112,7 @@ def dijkstra(
 
             heapq.heappush(queue, (dis + cost, neighbor, path + (neighbor,)))
 
+    if target is None:
+        return preds
+
     return preds.get(target, None)