improved typing and simplified track nodes

Author: BergJulian

yaramo/edge.py

@@ -1,10 +1,15 @@
-from typing import List, Optional, Tuple
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from yaramo.base_element import BaseElement
 from yaramo.geo_node import GeoNode
 from yaramo.node import Node
 from yaramo.vacancy_section import VacancySection
 
+if TYPE_CHECKING:
+    from yaramo.signal import Signal, SignalDirection
+
 
 class Edge(BaseElement):
     """This class is one of two Elements (Edge and Node) comprising the base of the yaramo Topology.
@@ -22,7 +27,7 @@ def __init__(
         vacancy_section: Optional[VacancySection] = None,
         length: float = None,
         intermediate_geo_nodes: List[GeoNode] = None,
-        signals: list["Signal"] = None,
+        signals: list[Signal] = None,
         maximum_speed: int = None,
         **kwargs
     ):
@@ -86,7 +91,7 @@ def __get_length(self) -> float:
         )
         return total_length
 
-    def get_direction_based_on_start_node(self, start: "Node") -> "SignalDirection":
+    def get_direction_based_on_start_node(self, start: Node) -> SignalDirection:
         from yaramo.signal import SignalDirection
 
         if self.node_a.uuid == start.uuid:
@@ -95,7 +100,7 @@ def get_direction_based_on_start_node(self, start: "Node") -> "SignalDirection":
             return SignalDirection.GEGEN
         return None
 
-    def get_direction_based_on_nodes(self, node_a: "Node", node_b: "Node") -> "SignalDirection":
+    def get_direction_based_on_nodes(self, node_a: Node, node_b: Node) -> SignalDirection:
         """Returns the direction according to whether the order of node_a and node_b is the same as in self
 
         Parameters
@@ -119,7 +124,7 @@ def get_direction_based_on_nodes(self, node_a: "Node", node_b: "Node") -> "Signa
             return SignalDirection.GEGEN
         return None
 
-    def get_signals_with_direction_in_order(self, direction: "SignalDirection") -> List["Signal"]:
+    def get_signals_with_direction_in_order(self, direction: SignalDirection) -> List[Signal]:
         """Returns all the signals (with that direction) on that Edge ordered by the given direction
 
         This only consideres Signals of SignalFunction type Einfahr_Signal, Ausfahr_Signal and Block_Signal that have the same direction as requested.
@@ -150,7 +155,7 @@ def get_signals_with_direction_in_order(self, direction: "SignalDirection") -> L
         result.sort(key=lambda x: x.distance_edge, reverse=(direction == SignalDirection.GEGEN))
         return result
 
-    def get_opposite_node(self, node: "Node") -> "Node":
+    def get_opposite_node(self, node: Node) -> Node:
         """Returns the opposite Node of the given Node
 
         Parameters
@@ -168,7 +173,7 @@ def get_opposite_node(self, node: "Node") -> "Node":
             return self.node_b
         return self.node_a
 
-    def get_next_geo_node(self, node: "Node") -> "GeoNode":
+    def get_next_geo_node(self, node: Node) -> GeoNode:
         """Returns the next GeoNode on Edgeof the given Top Node
 
         Parameters

yaramo/signal.py

@@ -1,5 +1,5 @@
-import logging
 from enum import Enum, auto
+import logging
 from typing import Set, Tuple
 from uuid import uuid4
 

yaramo/track.py

@@ -33,6 +33,10 @@ def add_edge_section(self, edge: Edge, section_start: float, section_end: float)
     def edges(self) -> List[Edge]:
         return list(self.edge_sections.keys())
 
+    @property
+    def nodes(self) -> List[Node]:
+        return list({node for edge in self.edges for node in (edge.node_a, edge.node_b)})
+
     def get_edges_in_order(self) -> List[Edge]:
         if len(self.edges) <= 1:
             return self.edges
@@ -83,24 +87,3 @@ def is_node_in_track(self, node: Node):
             if edge.is_node_connected(node):
                 return True
         return False
-
-    def get_nodes_in_order(self) -> List[Node]:
-        nodes = [node for edge in self.edges for node in (edge.node_a, edge.node_b)]
-        start_nodes = [node for node, count in Counter(nodes).items() if count == 1]
-
-        if len(start_nodes) != 2:
-            raise ValueError("Edges of Track separated")
-
-        previous_node = start_nodes[0]
-        start_edge = next(edge for edge in self.edges if previous_node in (edge.node_a, edge.node_b))
-
-        edges_in_order = self.get_edges_in_order()
-        if edges_in_order[0] != start_edge:
-            edges_in_order.reverse()
-
-        nodes_in_order: list[Node] = [previous_node]
-        for edge in edges_in_order:
-            previous_node = edge.get_other_node(previous_node)
-            nodes_in_order.append(previous_node)
-
-        return nodes_in_order