diff --git a/src/editor.py b/src/editor.py index 234d9da..fc75b99 100644 --- a/src/editor.py +++ b/src/editor.py @@ -47,6 +47,7 @@ class Editor: self.node_radius: int = 10 self.line_size: int = (int) (self.node_radius / 5) self.selected_node: int = -1 + self.selected_edge: int = -1 def mainloop(self) -> None: self.state = State.LOADING @@ -256,6 +257,9 @@ class Editor: for edge in self.graph.edges: node_1, node_2 = self.graph.get_edge_nodes(edge) pygame.draw.line(self.win, (255, 0, 0), self.world_to_screen(node_1.x, node_1.z), self.world_to_screen(node_2.x, node_2.z), self.line_size) + if self.selected_edge != -1: + node_1, node_2 = self.graph.get_edge_nodes(self.graph.edges[self.selected_edge]) + pygame.draw.line(self.win, (0, 255, 255), self.world_to_screen(node_1.x, node_1.z), self.world_to_screen(node_2.x, node_2.z), self.line_size) def render_nodes(self) -> None: for node in self.graph.nodes: @@ -303,10 +307,17 @@ class Editor: return int(screen_x), int(screen_y) def select_node(self, node: int): + self.selected_edge = -1 if self.selected_node != -1 and node != -1: - self.create_edge(self.selected_node, node) + self.link_nodes(self.selected_node, node) self.selected_node = node + def link_nodes(self, node_1: int, node_2: int): + if self.graph.edge_exists(node_1, node_2): + self.selected_edge = self.graph.get_edge(node_1, node_2) + else: + self.create_edge(node_1, node_2) + def create_node(self) -> None: self.graph.add_node(self.node_candidate_pos[0], self.node_candidate_pos[1], self.typing_text) self.typing_text = "" diff --git a/src/graph/graph.py b/src/graph/graph.py index 0844f80..cfeb64f 100644 --- a/src/graph/graph.py +++ b/src/graph/graph.py @@ -18,11 +18,20 @@ class Graph: def number_of_nodes(self) -> int: return len(self.nodes) - def get_edge_nodes(self, edge) -> tuple[Node, Node]: + def get_edge(self, node_1: int, node_2: int) -> int: + for edge in self.edges: + if (edge.start == node_1 and edge.end == node_2) or (edge.start == node_2 and edge.end == node_1): + return self.edges.index(edge) + return -1 + + def get_edge_nodes(self, edge: Edge) -> tuple[Node, Node]: return self.nodes[edge.start], self.nodes[edge.end] def edges_adjacent_to(self, node_i: int) -> Iterator[Edge]: return filter(lambda e: e.start == node_i or e.end == node_i, self.edges) + + def edge_exists(self, node_1: int, node_2: int) -> bool: + return self.get_edge(node_1, node_2) != -1 def dijkstra(self, source_index: int, target_index: int) -> Optional[list[int]]: n = len(self.nodes)