Modernize and fix Python graph examples with type hints, docstrings, and improvements

Co-authored-by: djeada <37275728+djeada@users.noreply.github.com>

Author: Copilot

src/graphs/python/a_star/src/a_star.py

@@ -1,50 +1,73 @@
+"""A* search algorithm implementation."""
+
+from __future__ import annotations
+
+import math
 from heapq import heappush, heappop
+from typing import Callable, Any
+
+from graph import Graph, Vertex
 
 
-def a_star(graph, source, destination, heuristic):
+def a_star(
+    graph: Graph,
+    source: Vertex,
+    destination: Vertex,
+    heuristic: Callable[[Any, Any], float],
+) -> float:
+    """
+    Find the shortest path between source and destination using A* algorithm.
 
+    Args:
+        graph: The graph to search.
+        source: The starting vertex.
+        destination: The target vertex.
+        heuristic: A function that estimates the distance between two vertex values.
+
+    Returns:
+        The shortest distance from source to destination, or infinity if no path exists.
+    """
     if not (graph.contains(source) and graph.contains(destination)):
-        return float("inf")
+        return math.inf
 
-    distance_from_source = dict()
-    distance_till_destination = dict()
-    # previous_vertex = dict()
+    distance_from_source: dict[Vertex, float] = {}
+    distance_till_destination: dict[Vertex, float] = {}
 
     for vertex in graph.vertices():
-        distance_from_source[vertex] = float("inf")
-        distance_till_destination[vertex] = float("inf")
+        distance_from_source[vertex] = math.inf
+        distance_till_destination[vertex] = math.inf
 
     distance_from_source[source] = 0
     distance_till_destination[source] = heuristic(source.value, destination.value)
 
-    visited = list()
-    heappush(visited, (distance_till_destination[source], source))
-
-    while visited:
+    open_set: list[tuple[float, Vertex]] = []
+    heappush(open_set, (distance_till_destination[source], source))
+    in_open_set: set[Vertex] = {source}
 
-        current = heappop(visited)[1]
+    while open_set:
+        current = heappop(open_set)[1]
+        in_open_set.discard(current)
 
         if current == destination:
             return distance_from_source[destination]
-            # return reconstruct_path(previous_vertex, current)
 
         for edge in graph.edges_from_vertex(current):
-
             tentative_distance_from_source = (
                 distance_from_source[current] + edge.distance
             )
 
             if tentative_distance_from_source < distance_from_source[edge.destination]:
-                # previous_vertex[edge.destination] = current
                 distance_from_source[edge.destination] = tentative_distance_from_source
-                distance_till_destination[edge.destination] = distance_from_source[
-                    edge.destination
-                ] + heuristic(edge.destination.value, destination.value)
+                distance_till_destination[edge.destination] = (
+                    distance_from_source[edge.destination]
+                    + heuristic(edge.destination.value, destination.value)
+                )
 
-                if edge.destination not in visited:
+                if edge.destination not in in_open_set:
                     heappush(
-                        visited,
+                        open_set,
                         (distance_till_destination[edge.destination], edge.destination),
                     )
+                    in_open_set.add(edge.destination)
 
-    return float("inf")
+    return math.inf

src/graphs/python/a_star/src/graph.py

@@ -1,89 +1,102 @@
+"""Graph data structures for graph algorithms."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
 class Vertex:
-    def __init__(self, value):
-        self.value = value
+    """A vertex in a graph with an associated value."""
 
-    def __eq__(self, other):
+    def __init__(self, value: T) -> None:
+        self.value = value
 
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Vertex):
             return self.value == other.value
-
         return False
 
-    def __lt__(self, other):
-
+    def __lt__(self, other: object) -> bool:
         if not isinstance(other, Vertex):
             raise NotImplementedError
-
         return self.value < other.value
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return str(self.value)
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(self.value)
 
 
 class Edge:
-    def __init__(self, source, destination, distance):
+    """A weighted directed edge connecting two vertices."""
+
+    def __init__(self, source: Vertex, destination: Vertex, distance: float) -> None:
         self.source = source
         self.destination = destination
         self.distance = distance
 
-    def __eq__(self, other):
-
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, Edge):
             return (
                 self.source == other.source
                 and self.destination == other.destination
                 and self.distance == other.distance
             )
-
         return False
 
-    def __lt__(self, other):
-
+    def __lt__(self, other: object) -> bool:
         if not isinstance(other, Edge):
             raise NotImplementedError
-
         return self.distance < other.distance
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.source}, {self.destination}, {self.distance})"
 
 
 class Graph:
-    def __init__(self):
-        self._adj_dict = dict()
+    """A directed weighted graph using adjacency list representation."""
 
-    def add_edge(self, edge):
+    def __init__(self) -> None:
+        self._adj_dict: dict[Vertex, list[Edge]] = {}
 
+    def add_edge(self, edge: Edge) -> None:
+        """Add an edge to the graph."""
         if edge.source in self._adj_dict:
             self._adj_dict[edge.source].append(edge)
-
         else:
             self._adj_dict[edge.source] = [edge]
 
         if edge.destination not in self._adj_dict:
-            self._adj_dict[edge.destination] = list()
+            self._adj_dict[edge.destination] = []
 
-    def add_vertex(self, vertex):
+    def add_vertex(self, vertex: Vertex) -> None:
+        """Add a vertex to the graph."""
         if vertex not in self._adj_dict:
-            self._adj_dict[vertex] = list()
+            self._adj_dict[vertex] = []
 
-    def vertices(self):
+    def vertices(self) -> list[Vertex]:
+        """Return a list of all vertices in the graph."""
         return list(self._adj_dict.keys())
 
-    def edges(self):
+    def edges(self) -> list[Edge]:
+        """Return a list of all edges in the graph."""
         return [elem for array in self._adj_dict.values() for elem in array if array]
 
-    def edges_from_vertex(self, vertex):
+    def edges_from_vertex(self, vertex: Vertex) -> list[Edge]:
+        """Return all edges originating from the given vertex."""
         self._assert_graph_contains_vertex(vertex)
         return self._adj_dict[vertex]
 
-    def contains(self, vertex):
+    def contains(self, vertex: Vertex) -> bool:
+        """Check if the graph contains the given vertex."""
         return vertex in self._adj_dict
 
-    def connected(self, vertex_a, vertex_b):
+    def connected(self, vertex_a: Vertex, vertex_b: Vertex) -> bool:
+        """Check if two vertices are directly connected by an edge."""
         self._assert_graph_contains_vertex(vertex_a)
         self._assert_graph_contains_vertex(vertex_b)
 
@@ -97,23 +110,23 @@ def connected(self, vertex_a, vertex_b):
 
         return False
 
-    def empty(self):
+    def empty(self) -> bool:
+        """Check if the graph has no vertices."""
         return len(self._adj_dict) == 0
 
-    def size(self):
+    def size(self) -> int:
+        """Return the number of vertices in the graph."""
         return len(self._adj_dict)
 
-    def __repr__(self):
-
+    def __repr__(self) -> str:
         result = ""
         for vertex, edges in self._adj_dict.items():
             result += f"{vertex} —> "
             for edge in edges:
-                result += edge
+                result += str(edge)
             result += "\n"
-
         return result
 
-    def _assert_graph_contains_vertex(self, vertex):
+    def _assert_graph_contains_vertex(self, vertex: Vertex) -> None:
         if vertex not in self._adj_dict:
-            raise KeyError(f"Graph doesn't containt the {vertex} vertex!")
+            raise KeyError(f"Graph doesn't contain the {vertex} vertex!")

src/graphs/python/a_star/tests/test_a_star.py

@@ -42,10 +42,10 @@ def distance(point_a, point_b):
         self.assertEqual(a_star(graph, Vertex(a), Vertex(c), Point.distance), 18)
         self.assertEqual(a_star(graph, Vertex(a), Vertex(b), Point.distance), 6)
         self.assertEqual(
-            a_star(graph, Vertex(d), Vertex(b), Point.distance), float("inf")
+            a_star(graph, Vertex(d), Vertex(b), Point.distance), math.inf
         )
         self.assertEqual(
-            a_star(graph, Vertex(b), Vertex(e), Point.distance), float("inf")
+            a_star(graph, Vertex(b), Vertex(e), Point.distance), math.inf
         )
 
 

src/graphs/python/bellman_ford/src/bellman_ford.py

@@ -1,40 +1,56 @@
-from graph import Graph
+"""Bellman-Ford algorithm for shortest paths with negative edges."""
 
+from __future__ import annotations
 
-def bellman_ford(graph, source, destination):
+import math
 
+from graph import Graph, Vertex
+
+
+def bellman_ford(graph: Graph, source: Vertex, destination: Vertex) -> float:
+    """
+    Find the shortest path between source and destination using Bellman-Ford algorithm.
+
+    This algorithm handles negative edge weights and detects negative cycles.
+
+    Args:
+        graph: The graph to search.
+        source: The starting vertex.
+        destination: The target vertex.
+
+    Returns:
+        The shortest distance from source to destination, or infinity if no path exists
+        or if a negative cycle is detected.
+    """
     if not (graph.contains(source) and graph.contains(destination)):
-        return float("inf")
+        return math.inf
 
     if source == destination:
         return 0
 
-    distances = dict()
-    visited = dict()
+    distances: dict[Vertex, float] = {}
 
     for vertex in graph.vertices():
-        distances[vertex] = float("inf")
-        visited[vertex] = False
+        distances[vertex] = math.inf
 
     distances[source] = 0
 
-    for i in range(graph.size() - 1):
+    for _ in range(graph.size() - 1):
         for edge in graph.edges():
-
             new_distance = distances[edge.source] + edge.distance
 
             if (
-                distances[edge.source] != float("Inf")
+                distances[edge.source] != math.inf
                 and new_distance < distances[edge.destination]
             ):
                 distances[edge.destination] = new_distance
 
     for edge in graph.edges():
         if (
-            distances[edge.source] != float("Inf")
+            distances[edge.source] != math.inf
             and distances[edge.source] + edge.distance < distances[edge.destination]
         ):
             # Graph contains negative weight cycle
-            return float("inf")
+            return math.inf
 
     return distances[destination]