diff --git a/graphs/a_star.py b/graphs/a_star.py
index d3657cb19..e0f24734a 100644
--- a/graphs/a_star.py
+++ b/graphs/a_star.py
@@ -1,37 +1,21 @@
-grid = [
-    [0, 1, 0, 0, 0, 0],
-    [0, 1, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
-    [0, 1, 0, 0, 0, 0],
-    [0, 1, 0, 0, 1, 0],
-    [0, 0, 0, 0, 1, 0],
+from __future__ import annotations
+
+DIRECTIONS = [
+    [-1, 0],  # left
+    [0, -1],  # down
+    [1, 0],  # right
+    [0, 1],  # up
 ]
 
-"""
-heuristic = [[9, 8, 7, 6, 5, 4],
-             [8, 7, 6, 5, 4, 3],
-             [7, 6, 5, 4, 3, 2],
-             [6, 5, 4, 3, 2, 1],
-             [5, 4, 3, 2, 1, 0]]"""
-
-init = [0, 0]
-goal = [len(grid) - 1, len(grid[0]) - 1]  # all coordinates are given in format [y,x]
-cost = 1
-
-# the cost map which pushes the path closer to the goal
-heuristic = [[0 for row in range(len(grid[0]))] for col in range(len(grid))]
-for i in range(len(grid)):
-    for j in range(len(grid[0])):
-        heuristic[i][j] = abs(i - goal[0]) + abs(j - goal[1])
-        if grid[i][j] == 1:
-            heuristic[i][j] = 99  # added extra penalty in the heuristic map
-
-
-# the actions we can take
-delta = [[-1, 0], [0, -1], [1, 0], [0, 1]]  # go up  # go left  # go down  # go right
-
 
 # function to search the path
-def search(grid, init, goal, cost, heuristic):
+def search(
+    grid: list[list[int]],
+    init: list[int],
+    goal: list[int],
+    cost: int,
+    heuristic: list[list[int]],
+) -> tuple[list[list[int]], list[list[int]]]:
 
     closed = [
         [0 for col in range(len(grid[0]))] for row in range(len(grid))
@@ -52,7 +36,7 @@ def search(grid, init, goal, cost, heuristic):
 
     while not found and not resign:
         if len(cell) == 0:
-            return "FAIL"
+            raise ValueError("Algorithm is unable to find solution")
         else:  # to choose the least costliest action so as to move closer to the goal
             cell.sort()
             cell.reverse()
@@ -64,9 +48,9 @@ def search(grid, init, goal, cost, heuristic):
             if x == goal[0] and y == goal[1]:
                 found = True
             else:
-                for i in range(len(delta)):  # to try out different valid actions
-                    x2 = x + delta[i][0]
-                    y2 = y + delta[i][1]
+                for i in range(len(DIRECTIONS)):  # to try out different valid actions
+                    x2 = x + DIRECTIONS[i][0]
+                    y2 = y + DIRECTIONS[i][1]
                     if x2 >= 0 and x2 < len(grid) and y2 >= 0 and y2 < len(grid[0]):
                         if closed[x2][y2] == 0 and grid[x2][y2] == 0:
                             g2 = g + cost
@@ -79,8 +63,8 @@ def search(grid, init, goal, cost, heuristic):
     y = goal[1]
     invpath.append([x, y])  # we get the reverse path from here
     while x != init[0] or y != init[1]:
-        x2 = x - delta[action[x][y]][0]
-        y2 = y - delta[action[x][y]][1]
+        x2 = x - DIRECTIONS[action[x][y]][0]
+        y2 = y - DIRECTIONS[action[x][y]][1]
         x = x2
         y = y2
         invpath.append([x, y])
@@ -88,13 +72,37 @@ def search(grid, init, goal, cost, heuristic):
     path = []
     for i in range(len(invpath)):
         path.append(invpath[len(invpath) - 1 - i])
+    return path, action
+
+
+if __name__ == "__main__":
+    grid = [
+        [0, 1, 0, 0, 0, 0],
+        [0, 1, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
+        [0, 1, 0, 0, 0, 0],
+        [0, 1, 0, 0, 1, 0],
+        [0, 0, 0, 0, 1, 0],
+    ]
+
+    init = [0, 0]
+    # all coordinates are given in format [y,x]
+    goal = [len(grid) - 1, len(grid[0]) - 1]
+    cost = 1
+
+    # the cost map which pushes the path closer to the goal
+    heuristic = [[0 for row in range(len(grid[0]))] for col in range(len(grid))]
+    for i in range(len(grid)):
+        for j in range(len(grid[0])):
+            heuristic[i][j] = abs(i - goal[0]) + abs(j - goal[1])
+            if grid[i][j] == 1:
+                # added extra penalty in the heuristic map
+                heuristic[i][j] = 99
+
+    path, action = search(grid, init, goal, cost, heuristic)
+
     print("ACTION MAP")
     for i in range(len(action)):
         print(action[i])
 
-    return path
-
-
-a = search(grid, init, goal, cost, heuristic)
-for i in range(len(a)):
-    print(a[i])
+    for i in range(len(path)):
+        print(path[i])