fix geo nodes for edge

simulate-digital-rail · Jan 29, 2024 · 107f382 · 107f382
1 parent ef57d08
commit 107f382
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/yaramo/edge.py b/yaramo/edge.py
@@ -16,7 +16,7 @@ class Edge(BaseElement):
     """
 
     def __init__(self, node_a: Node, node_b: Node, vacancy_section: Optional[VacancySection] = None, length: float = None,
-                 intermediate_geo_nodes=None, **kwargs):
+                 intermediate_geo_nodes: List[GeoNode]=None, signals: list["Signal"]=None, maximum_speed: int=None, **kwargs):
         """
         Parameters
         ----------
@@ -34,10 +34,9 @@ def __init__(self, node_a: Node, node_b: Node, vacancy_section: Optional[Vacancy
         self.intermediate_geo_nodes = intermediate_geo_nodes or []
         self.node_a = node_a
         self.node_b = node_b
-        self.intermediate_geo_nodes: list[GeoNode] = []
-        self.signals: list[Signal] = []
+        self.signals = signals or []
         self.length = length
-        self.maximum_speed: int = None
+        self.maximum_speed = maximum_speed
         self.vacancy_section = vacancy_section
 
     def is_node_connected(self, other_node) -> bool:

diff --git a/yaramo/topology.py b/yaramo/topology.py
@@ -2,6 +2,7 @@
 
 from yaramo.base_element import BaseElement
 from yaramo.edge import Edge
+from yaramo.geo_node import Wgs84GeoNode
 from yaramo.node import Node
 from yaramo.route import Route
 from yaramo.signal import Signal
@@ -92,10 +93,11 @@ def from_json(cls, json_str: str):
             node_b = topology.nodes[edge["node_b"]]
             node_a.connected_nodes.append(node_b)
             node_b.connected_nodes.append(node_a)
+            geo_node_strings = [obj["objects"][geo_node_uuid] for geo_node_uuid in edge["intermediate_geo_nodes"] if geo_node_uuid in obj["objects"]]
             topology.add_edge(Edge(**{
                 **edge, "node_a": node_a, "node_b": node_b,
                 "signals": [topology.signals[signal_uuid] for signal_uuid in edge["signals"]],
-                "intermediate_geo_nodes": [obj["objects"][geo_node_uuid] for geo_node_uuid in edge["intermediate_geo_nodes"] if geo_node_uuid in obj["objects"]],
+                "intermediate_geo_nodes": [Wgs84GeoNode(obj["objects"][geo_node["geo_point"]]["x"], obj["objects"][geo_node["geo_point"]]["y"], name=geo_node["name"], uuid=geo_node["uuid"]) for geo_node in geo_node_strings]
             }))
         for signal in obj["signals"]:
             topology.signals[signal["uuid"]].edge = topology.edges[signal["edge"]]