add doc string

metadrive/manager/sumo_map_manager.py

@@ -1,6 +1,6 @@
 from metadrive.component.map.scenario_map import ScenarioMap
 from metadrive.manager.base_manager import BaseManager
-from metadrive.utils.sumo.map_utils import StreetMap, extract_map_features
+from metadrive.utils.sumo.map_utils import extract_map_features, RoadLaneJunctionGraph
 
 
 class SumoMapManager(BaseManager):
@@ -10,22 +10,33 @@ class SumoMapManager(BaseManager):
     PRIORITY = 0  # Map update has the most high priority
 
     def __init__(self, sumo_map_path):
+        """
+        Init the map manager. It can be extended to manage more maps
+        """
         super(SumoMapManager, self).__init__()
         self.current_map = None
-        street_map = StreetMap()
-        street_map.reset(sumo_map_path)
-        self.map_feature = extract_map_features(street_map)
+        self.graph = RoadLaneJunctionGraph(sumo_map_path)
+        self.map_feature = extract_map_features(self.graph)
 
     def destroy(self):
+        """
+        Delete the map manager
+        """
         self.current_map.destroy()
         super(SumoMapManager, self).destroy()
         self.current_map = None
 
     def before_reset(self):
+        """
+        Detach existing maps
+        """
         if self.current_map:
             self.current_map.detach_from_world()
 
     def reset(self):
+        """
+        Rebuild the map and load it into the scene
+        """
         if not self.current_map:
             self.current_map = ScenarioMap(map_index=0, map_data=self.map_feature)
         self.current_map.attach_to_world()

metadrive/tests/vis_env/vis_sumo_map.py

@@ -9,18 +9,23 @@
 
 class MyEnv(BaseEnv):
     def reward_function(self, agent):
+        """Dummy reward function."""
         return 0, {}
 
     def cost_function(self, agent):
+        """Dummy cost function."""
         return 0, {}
 
     def done_function(self, agent):
+        """Dummy done function."""
         return False, {}
 
     def get_single_observation(self):
+        """Dummy observation function."""
         return DummyObservation()
 
     def setup_engine(self):
+        """Register the map manager"""
         super().setup_engine()
         map_path = AssetLoader.file_path("carla", "CARLA_town01.net.xml", unix_style=False)
         self.engine.register_manager("map_manager", SumoMapManager(map_path))

metadrive/utils/sumo/map_utils.py

@@ -43,6 +43,9 @@ def __init__(
         shape,
         width: float,
     ):
+        """
+        Lane shape
+        """
         shape = buffered_shape(shape.getShape(), shape.getWidth())
         self.shape = shape
 
@@ -53,11 +56,16 @@ class RoadShape:
     right_border: np.ndarray
 
     def __post_init__(self):
+        """
+        post process
+        """
         self.polygon = self.left_border + list(reversed(self.right_border))
 
 
 class JunctionNode:
+
     def __init__(self, sumolib_obj):
+        """Node for junction node."""
         self.sumolib_obj: sumolib.net.node = sumolib_obj
         self.name = sumolib_obj.getID()
         self.type = sumolib_obj.getType()
@@ -73,6 +81,9 @@ def __init__(self, sumolib_obj):
 
 class LaneNode:
     def __init__(self, sumolib_obj):
+        """
+        Node for a lane
+        """
         self.sumolib_obj: sumolib.net.lane = sumolib_obj
         self.name: str = sumolib_obj.getID()
         self.edge_type: str = sumolib_obj.getEdge().getType()
@@ -111,6 +122,9 @@ def __init__(
         from_junction,
         to_junction,
     ):
+        """
+        Node for a road
+        """
         self.sumolib_obj: sumolib.net.edge = sumolib_obj
         self.name: str = sumolib_obj.getID()
         self.type = sumolib_obj.getType()
@@ -147,6 +161,7 @@ def __init__(
         self,
         sumo_net_path,
     ):
+        """Init the graph"""
 
         self.sumo_net = sumolib.net.readNet(
             sumo_net_path, withInternal=True, withPedestrianConnections=True, withPrograms=True
@@ -247,6 +262,7 @@ def __init__(
         self.edge_dividers = edge_dividers
 
     def _compute_traffic_dividers(self, threshold=1):
+        """Find the road dividers"""
         lane_dividers = []  # divider between lanes with same traffic direction
         edge_dividers = []  # divider between lanes with opposite traffic direction
         edge_borders = []
@@ -283,23 +299,15 @@ def _compute_traffic_dividers(self, threshold=1):
         return lane_dividers, edge_dividers
 
 
-class StreetMap:
-    def __init__(self):
-        self.graph = None
-
-    def reset(self, sumo_net_path: str):
-        self.sumo_net_path = sumo_net_path
-        self.graph = RoadLaneJunctionGraph(self.sumo_net_path)
-
-
-def extract_map_features(street_map: StreetMap):
+def extract_map_features(graph):
+    """This func extracts the map features like lanes/lanelines from the SUMO map"""
     from shapely.geometry import Polygon
 
     ret = {}
     # # build map boundary
     polygons = []
 
-    for junction_id, junction in street_map.graph.junctions.items():
+    for junction_id, junction in graph.junctions.items():
         if len(junction.shape) <= 2:
             continue
         boundary_polygon = Polygon(junction.shape)
@@ -312,7 +320,7 @@ def extract_map_features(street_map: StreetMap):
         }
 
     # build map lanes
-    for road_id, road in street_map.graph.roads.items():
+    for road_id, road in graph.roads.items():
         for lane in road.lanes:
 
             id = "lane_{}".format(lane.name)
@@ -341,11 +349,11 @@ def extract_map_features(street_map: StreetMap):
                     SD.POLYGON: boundary_polygon,
                 }
 
-    for lane_divider_id, lane_divider in enumerate(street_map.graph.lane_dividers):
+    for lane_divider_id, lane_divider in enumerate(graph.lane_dividers):
         id = "lane_divider_{}".format(lane_divider_id)
         ret[id] = {SD.TYPE: MetaDriveType.LINE_BROKEN_SINGLE_WHITE, SD.POLYLINE: lane_divider}
 
-    for edge_divider_id, edge_divider in enumerate(street_map.graph.edge_dividers):
+    for edge_divider_id, edge_divider in enumerate(graph.edge_dividers):
         id = "edge_divider_{}".format(edge_divider_id)
         ret[id] = {SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_YELLOW, SD.POLYLINE: edge_divider}