add rtree to graph itself

evrouting/osm.py

@@ -37,25 +37,11 @@ def query_url(service, coordinates, osrm_config: OsrmConf):
            f'{";".join([f"{lon},{lat}" for lat, lon in coordinates])}'
 
 
-def insert_charging_stations(G, charging_stations, osrm_conf=None):
-    osrm_conf = osrm_conf or OsrmConf(
-        server='localhost',
-        port=5000
-    )
-
-    index = rtree.index.Index()
-
-    for n in G.nodes:
-        lon = G.nodes[n]['lon']
-        lat = G.nodes[n]['lat']
-
-        # insert point
-        index.insert(n, (lon, lat, lon, lat))
-
+def insert_charging_stations(G, charging_stations):
     for s in charging_stations:
         lon = s['lon']
         lat = s['lat']
-        n = list(index.nearest((lon, lat, lon, lat), 1))[0]
+        n = list(G.rtree.nearest((lon, lat, lon, lat), 1))[0]
         G.nodes[n][CHARGING_COEFFICIENT_KEY] = s['power']
 
     return G
@@ -136,12 +122,15 @@ def read_osm(osm_xml_data, only_roads=True) -> nx.DiGraph:
             nx.add_path(G, w.nds, id=w.id)
             nx.add_path(G, w.nds[::-1], id=w.id)
 
+    G.rtree = rtree.index.Index()
+
     # Complete the used nodes' information
     for n_id in G.nodes():
         n = osm.nodes[n_id]
         G.nodes[n_id]['lat'] = n.lat
         G.nodes[n_id]['lon'] = n.lon
         G.nodes[n_id]['id'] = n.id
+        G.rtree.insert(n_id, (n.lat, n.lon, n.lat, n.lon))
 
     return G
 

tests/osm/test_osm_charge.py

 import os
 
+import pytest
 import networkx as nx
+import rtree
 
 from evrouting.osm import read_osm, insert_charging_stations
 from evrouting.graph_tools import CHARGING_COEFFICIENT_KEY
 
 
+@pytest.fixture
+def graph():
+    G = nx.DiGraph()
+    G.rtree = rtree.index.Index()
+
+    node_coordinates = [
+        (51.7705832, 7.0002595),
+        (51.7696529, 6.9568520)
+    ]
+
+    for n_id, coordinates in enumerate(node_coordinates):
+        lat, lon = coordinates
+        # Add two nodes, that exist in osm test map
+        G.add_node(n_id, lat=lat, lon=lon)
+        G.rtree.insert(n_id, (lon, lat, lon, lat))
+
+    yield G
+    del G
+
+
 def _test_read_osm():
     """Just check if it runs. Todo: Delete."""
     assert read_osm(os.path.join(os.path.dirname(__file__), 'static/map.osm'))
 
 
-def test_insert_charging_stations_close():
-    G = nx.DiGraph()
-
-    # Add two nodes, that exist in osm test map
-    G.add_node(0, lat=51.7705832, lon=7.0002595)
-    G.add_node(1, lat=51.7696529, lon=6.9568520)
-
+def test_insert_charging_stations_close(graph):
     # Close two node 1
     S = [{"lon": 7.0002593, "lat": 51.7705832, "power": 22.0}]
 
-    G = insert_charging_stations(G, S)
+    graph = insert_charging_stations(graph, S)
 
-    assert G.nodes[0][CHARGING_COEFFICIENT_KEY] == 22.0
-    assert CHARGING_COEFFICIENT_KEY not in G.nodes[1]
-
-
-def test_insert_charging_stations_eq():
-    G = nx.DiGraph()
+    assert graph.nodes[0][CHARGING_COEFFICIENT_KEY] == 22.0
+    assert CHARGING_COEFFICIENT_KEY not in graph.nodes[1]
 
-    # Add two nodes, that exist in osm test map
-    G.add_node(0, lat=51.7705832, lon=7.0002595)
-    G.add_node(1, lat=51.7696529, lon=6.9568520)
 
+def test_insert_charging_stations_eq(graph):
     # Close exactly at node 1
     S = [{"lon": 7.0002595, "lat": 51.7705832, "power": 22.0}]
 
-    G = insert_charging_stations(G, S)
+    graph = insert_charging_stations(graph, S)
 
-    assert G.nodes[0][CHARGING_COEFFICIENT_KEY] == 22.0
-    assert CHARGING_COEFFICIENT_KEY not in G.nodes[1]
+    assert graph.nodes[0][CHARGING_COEFFICIENT_KEY] == 22.0
+    assert CHARGING_COEFFICIENT_KEY not in graph.nodes[1]