From 122ec20f539fd8307f9ec3344279a54d6c76ac2f Mon Sep 17 00:00:00 2001
From: Hari <harisankar995@gmail.com>
Date: Fri, 19 Jan 2024 19:14:13 +0100
Subject: [PATCH] added additional helper functions

---
 pathfinding3d/core/heap.py | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/pathfinding3d/core/heap.py b/pathfinding3d/core/heap.py
index a98fe6d..78f91ca 100644
--- a/pathfinding3d/core/heap.py
+++ b/pathfinding3d/core/heap.py
@@ -30,6 +30,7 @@ def __init__(self, node: GridNode, grid: Union[Grid, World]):
 
         self.grid = grid
         self._get_node_tuple = self._determine_node_retrieval_function()
+        self._get_node = self._determine_node_function()
         self.open_list = [self._get_node_tuple(node, 0)]
         self.removed_node_tuples = set()
         self.heap_order = {}
@@ -57,6 +58,29 @@ def _determine_node_retrieval_function(self) -> Callable:
 
         raise ValueError("Unsupported grid type")
 
+    def _determine_node_function(self) -> Callable:
+        """
+        Determines the node function based on the type of grid.
+
+        Returns
+        -------
+        function
+            A function that takes a node tuple and returns the corresponding node.
+
+        Raises
+        ------
+        ValueError
+            If the grid is not of type Grid or World.
+        """
+
+        if isinstance(self.grid, Grid):
+            return lambda node_tuple: self.grid.node(*node_tuple[2:])
+
+        if isinstance(self.grid, World):
+            return lambda node_tuple: self.grid.grids[node_tuple[5]].node(*node_tuple[2:5])
+
+        raise ValueError("Unsupported grid type")
+
     def pop_node(self) -> GridNode:
         """
         Pops the node with the lowest cost from the heap.
@@ -70,12 +94,7 @@ def pop_node(self) -> GridNode:
         while node_tuple in self.removed_node_tuples:
             node_tuple = heapq.heappop(self.open_list)
 
-        if isinstance(self.grid, Grid):
-            node = self.grid.node(*node_tuple[2:])
-        elif isinstance(self.grid, World):
-            node = self.grid.grids[node_tuple[5]].node(*node_tuple[2:5])
-
-        return node
+        return self._get_node(node_tuple)
 
     def push_node(self, node: GridNode):
         """