typing

Author: BergJulian

yaramo/geo_node.py

@@ -1,7 +1,8 @@
-import math
-from abc import ABC, abstractmethod
+from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from haversine import Unit, haversine
+import math
 
 from yaramo.base_element import BaseElement
 
@@ -21,27 +22,27 @@ def __init__(self, x, y, dbref_crs: str = "ER0", **kwargs):
         self.dbref_crs = dbref_crs
 
     @abstractmethod
-    def get_distance_to_other_geo_node(self, geo_node_b: "GeoNode"):
+    def get_distance_to_other_geo_node(self, geo_node_b: GeoNode):
         """Returns to distance to the given other GeoNode."""
         pass
 
     @abstractmethod
-    def to_wgs84(self) -> "Wgs84GeoNode":
+    def to_wgs84(self) -> Wgs84GeoNode:
         pass
 
     @abstractmethod
-    def to_dbref(self) -> "DbrefGeoNode":
+    def to_dbref(self) -> DbrefGeoNode:
         pass
 
     @abstractmethod
-    def to_euclidean(self) -> "EuclideanGeoNode":
+    def to_euclidean(self) -> EuclideanGeoNode:
         pass
 
     def to_serializable(self):
         return self.__dict__, {}
 
     @staticmethod
-    def get_new_geo_node_same_type(old_geo_node: "GeoNode", x: float, y: float) -> "GeoNode":
+    def get_new_geo_node_same_type(old_geo_node: GeoNode, x: float, y: float) -> GeoNode:
         if isinstance(old_geo_node, Wgs84GeoNode):
             return Wgs84GeoNode(x, y)
         elif isinstance(old_geo_node, DbrefGeoNode):
@@ -52,39 +53,39 @@ def get_new_geo_node_same_type(old_geo_node: "GeoNode", x: float, y: float) -> "
 
 
 class Wgs84GeoNode(GeoNode):
-    def get_distance_to_other_geo_node(self, geo_node_b: "GeoNode"):
+    def get_distance_to_other_geo_node(self, geo_node_b: GeoNode):
         geo_node_b = geo_node_b.to_wgs84()
         return self.__haversine_distance(geo_node_b)
 
-    def __haversine_distance(self, geo_node_b: "GeoNode"):
+    def __haversine_distance(self, geo_node_b: GeoNode):
         own = (self.x, self.y)
         other = (geo_node_b.x, geo_node_b.y)
         return haversine(own, other, unit=Unit.METERS)
 
-    def to_wgs84(self) -> "Wgs84GeoNode":
+    def to_wgs84(self) -> Wgs84GeoNode:
         return self
 
-    def to_dbref(self) -> "DbrefGeoNode":
+    def to_dbref(self) -> DbrefGeoNode:
         x, y = transform_wgs84_to_dbref(self.x, self.y, self.dbref_crs)
         return DbrefGeoNode(x, y, self.dbref_crs, uuid=self.uuid)
 
-    def to_euclidean(self) -> "EuclideanGeoNode":
+    def to_euclidean(self) -> EuclideanGeoNode:
         return self.to_dbref().to_euclidean()
 
 
 class DbrefGeoNode(GeoNode):
-    def get_distance_to_other_geo_node(self, geo_node_b: "GeoNode"):
+    def get_distance_to_other_geo_node(self, geo_node_b: GeoNode):
         # Separate DB Ref distance method not implemented yet, therefore use WGS84 distance
         return self.to_wgs84().get_distance_to_other_geo_node(geo_node_b)
 
-    def to_wgs84(self) -> "Wgs84GeoNode":
+    def to_wgs84(self) -> Wgs84GeoNode:
         x, y = transform_dbref_to_wgs84(self.x, self.y, self.dbref_crs)
         return Wgs84GeoNode(x, y, self.dbref_crs, uuid=self.uuid)
 
-    def to_dbref(self) -> "DbrefGeoNode":
+    def to_dbref(self) -> DbrefGeoNode:
         return self
 
-    def to_euclidean(self) -> "EuclideanGeoNode":
+    def to_euclidean(self) -> EuclideanGeoNode:
         # This transformation is just for testing purposes and not correct, see documentation in EuclideanGeoNode.
         _x_shift = 4533770.0
         _y_shift = 5625780.0
@@ -100,24 +101,24 @@ class EuclideanGeoNode(GeoNode):
     DBRef geo nodes, but the conversion is just a simple coordinate shift and with this, probably incorrect.
     """
 
-    def get_distance_to_other_geo_node(self, geo_node_b: "EuclideanGeoNode"):
+    def get_distance_to_other_geo_node(self, geo_node_b: EuclideanGeoNode):
         geo_node_b = geo_node_b.to_euclidean()
         return self.__eucldian_distance(geo_node_b)
 
-    def __eucldian_distance(self, geo_node_b: "GeoNode"):
+    def __eucldian_distance(self, geo_node_b: GeoNode):
         min_x = min(self.x, geo_node_b.x)
         min_y = min(self.y, geo_node_b.y)
         max_x = max(self.x, geo_node_b.x)
         max_y = max(self.y, geo_node_b.y)
         return math.sqrt(math.pow(max_x - min_x, 2) + math.pow(max_y - min_y, 2))
 
-    def to_wgs84(self) -> "Wgs84GeoNode":
+    def to_wgs84(self) -> Wgs84GeoNode:
         return self.to_dbref().to_wgs84()
 
-    def to_dbref(self) -> "DbrefGeoNode":
+    def to_dbref(self) -> DbrefGeoNode:
         _x_shift = 4533770.0
         _y_shift = 5625780.0
         return DbrefGeoNode(self.x + _x_shift, self.y + _y_shift, self.dbref_crs, uuid=self.uuid)
 
-    def to_euclidean(self) -> "EuclideanGeoNode":
+    def to_euclidean(self) -> EuclideanGeoNode:
         return self

yaramo/node.py

@@ -1,12 +1,17 @@
+from __future__ import annotations
+
 import sys
 from enum import Enum
 from itertools import permutations
 from math import atan2, cos, sin
-from typing import List
+from typing import TYPE_CHECKING
 
 from yaramo.base_element import BaseElement
 from yaramo.geo_node import GeoNode
 
+if TYPE_CHECKING:
+    from yaramo.edge import Edge
+
 
 class EdgeConnectionDirection(Enum):
     Spitze = 0
@@ -33,17 +38,17 @@ def __init__(self, turnout_side=None, **kwargs):
         """
 
         super().__init__(**kwargs)
-        self.connected_edge_on_head = None
-        self.connected_edge_on_right = None
-        self.connected_edge_on_left = None
+        self.connected_edge_on_head: Edge = None
+        self.connected_edge_on_right: Edge = None
+        self.connected_edge_on_left: Edge = None
         self.maximum_speed_on_left = None
         self.maximum_speed_on_right = None
-        self.connected_edges: list["Edge"] = []
+        self.connected_edges: list[Edge] = []
         self.geo_node: GeoNode = kwargs.get("geo_node", None)
         self.turnout_side: str = turnout_side
         self.drive_amount = 0
 
-    def maximum_speed(self, node_a: "Node", node_b: "Node"):
+    def maximum_speed(self, node_a: Node, node_b: Node):
         """Return the maximum allowed speed for traversing this node,
         coming from node_a and going to node_b
         """
@@ -56,49 +61,49 @@ def maximum_speed(self, node_a: "Node", node_b: "Node"):
             return None
 
     @property
-    def connected_nodes(self):
+    def connected_nodes(self) -> list[Node]:
         return [edge.get_opposite_node(self) for edge in self.connected_edges]
 
     @property
-    def connected_on_head(self):
+    def connected_on_head(self) -> Node:
         if self.connected_edge_on_head is None:
             self.calc_anschluss_of_all_edges()
         if self.connected_edge_on_head is None:
             return None
         return self.connected_edge_on_head.get_opposite_node(self)
 
     @property
-    def connected_on_left(self):
+    def connected_on_left(self) -> Node:
         if self.connected_edge_on_head is None:
             self.calc_anschluss_of_all_edges()
         if self.connected_edge_on_left is None:
             return None
         return self.connected_edge_on_left.get_opposite_node(self)
 
     @property
-    def connected_on_right(self):
+    def connected_on_right(self) -> Node:
         if self.connected_edge_on_head is None:
             self.calc_anschluss_of_all_edges()
         if self.connected_edge_on_right is None:
             return None
         return self.connected_edge_on_right.get_opposite_node(self)
 
-    def set_connection_head_edge(self, edge: "Edge"):
+    def set_connection_head_edge(self, edge: Edge):
         self.connected_edge_on_head = edge
         if edge not in self.connected_edges:
             self.connected_edges.append(edge)
 
-    def set_connection_left_edge(self, edge: "Edge"):
+    def set_connection_left_edge(self, edge: Edge):
         self.connected_edge_on_left = edge
         if edge not in self.connected_edges:
             self.connected_edges.append(edge)
 
-    def set_connection_right_edge(self, edge: "Edge"):
+    def set_connection_right_edge(self, edge: Edge):
         self.connected_edge_on_right = edge
         if edge not in self.connected_edges:
             self.connected_edges.append(edge)
 
-    def remove_edge(self, edge: "Edge"):
+    def remove_edge(self, edge: Edge):
         self.connected_edges.remove(edge)
         if self.connected_edge_on_head == edge:
             self.connected_edge_on_head = None
@@ -107,7 +112,7 @@ def remove_edge(self, edge: "Edge"):
         if self.connected_edge_on_right == edge:
             self.connected_edge_on_right = None
 
-    def remove_edge_to_node(self, node: "Node"):
+    def remove_edge_to_node(self, node: Node):
         """Removes the edge to the given node and removes the node from the connected_nodes list."""
         edge = self.get_edge_to_node(node)
         self.remove_edge(edge)
@@ -116,7 +121,7 @@ def get_edge_to_node(self, node):
         """Returns the edge to the given neighbor node."""
         return next(edge for edge in self.connected_edges if edge.get_opposite_node(self) == node)
 
-    def get_possible_followers(self, source: "Edge") -> List["Edge"]:
+    def get_possible_followers(self, source: Edge) -> list[Edge]:
         """Returns the `Edge`s that could follow (head, left, right) when comming from a source `Edge` connected to this `Node`."""
         if source is None:
             return self.connected_edges
@@ -131,7 +136,7 @@ def get_possible_followers(self, source: "Edge") -> List["Edge"]:
             return [self.connected_edge_on_left, self.connected_edge_on_right]
         return [self.connected_edge_on_head]
 
-    def get_anschluss_for_edge(self, edge: "Edge") -> EdgeConnectionDirection:
+    def get_anschluss_for_edge(self, edge: Edge) -> EdgeConnectionDirection:
         """Gets the Anschluss (Ende, Links, Rechts, Spitze) of other node.
         Idea: We assume, the current node is a point and we want to estimate the Anschluss of the other node.
 
@@ -203,7 +208,7 @@ def get_rad_between_nodes(geo_node_a: GeoNode, geo_node_b: GeoNode) -> float:
             self.connected_edge_on_right = right
             break
 
-    def is_point(self):
+    def is_point(self) -> bool:
         """
         Returns true if this node is a point.
         A point is a `Node` with at least 3 connected tracks (yaramo only supports three edges at the moment)