pycityagent 2.0.0a65__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a67__cp310-cp310-macosx_11_0_arm64.whl

pycityagent/environment/simulator.py

@@ -4,6 +4,7 @@ import asyncio
 import logging
 import os
 from datetime import datetime, timedelta
+import time
 from typing import Any, Optional, Union, cast
 
 from mosstool.type import TripMode
@@ -20,11 +21,18 @@ from .utils.const import *
 
 logger = logging.getLogger("pycityagent")
 
+__all__ = [
+    "Simulator",
+]
+
 
 class Simulator:
     """
-    - 模拟器主类
-    - Simulator Class
+    Main class of the simulator.
+
+    - **Description**:
+        - This class is the core of the simulator, responsible for initializing and managing the simulation environment.
+        - It reads parameters from a configuration dictionary, initializes map data, and starts or connects to a simulation server as needed.
     """
 
     def __init__(self, config: dict, secure: bool = False) -> None:
@@ -106,33 +114,74 @@ class Simulator:
             poi["id"]: poi["aoi_id"] for _, poi in self.map.pois.items()
         }
         self._environment_prompt:dict[str, str] = {}
+        self._log_list = []
+
+    def get_log_list(self):
+        return self._log_list
+    
+    def clear_log_list(self):
+        self._log_list = []
 
     @property
-    def environment(self):
+    def environment(self) -> dict[str, str]:
+        """
+        Get the current state of environment variables.
+        """
         return self._environment_prompt
-    
+
     def set_environment(self, environment: dict[str, str]):
+        """
+        Set the entire dictionary of environment variables.
+
+        - **Args**:
+            - `environment` (`Dict[str, str]`): Key-value pairs of environment variables.
+        """
         self._environment_prompt = environment
 
-    def sence(self, key: str):
+    def sence(self, key: str) -> str:
+        """
+        Retrieve the value of an environment variable by its key.
+
+        - **Args**:
+            - `key` (`str`): The key of the environment variable.
+
+        - **Returns**:
+            - `str`: The value of the corresponding key, or an empty string if not found.
+        """
         return self._environment_prompt.get(key, "")
-    
+
     def update_environment(self, key: str, value: str):
+        """
+        Update the value of a single environment variable.
+
+        - **Args**:
+            - `key` (`str`): The key of the environment variable.
+            - `value` (`str`): The new value to set.
+        """
         self._environment_prompt[key] = value
 
     # * Agent相关
     def find_agents_by_area(self, req: dict, status=None):
         """
-        通过区域范围查找agent/person
-        Get agents/persons in the provided area
+        Find agents/persons within a specified area.
 
-        Args:
-        - req (dict): 用于描述区域的请求 https://cityproto.sim.fiblab.net/#city.person.1.GetPersonByLongLatBBoxRequest
-        - status (int): 用于限制agent/person状态 if 'status' is not None, then you get those persons in 'status' https://cityproto.sim.fiblab.net/#city.agent.v2.Status
+        - **Args**:
+            - `req` (`dict`): A dictionary that describes the area. Refer to
+              https://cityproto.sim.fiblab.net/#city.person.1.GetPersonByLongLatBBoxRequest.
+            - `status` (`Optional[int]`): An integer representing the status of the agents/persons to filter by.
+              If provided, only persons with the given status will be returned.
+              Refer to https://cityproto.sim.fiblab.net/#city.agent.v2.Status.
 
-        Returns:
-        - https://cityproto.sim.fiblab.net/#city.person.1.GetPersonByLongLatBBoxResponse
+        - **Returns**:
+            - The response from the GetPersonByLongLatBBox method, possibly filtered by status.
+              Refer to https://cityproto.sim.fiblab.net/#city.person.1.GetPersonByLongLatBBoxResponse.
         """
+        start_time = time.time()
+        log = {
+            "req": "find_agents_by_area",
+            "start_time": start_time,
+            "consumption": 0
+        }
         loop = asyncio.get_event_loop()
         resp = loop.run_until_complete(
             self._client.person_service.GetPersonByLongLatBBox(req=req)
@@ -146,6 +195,8 @@ class Simulator:
                 if agent.status in status:
                     motions.append(agent)
             resp.motions = motions  # type: ignore
+            log["consumption"] = time.time() - start_time
+            self._log_list.append(log)
             return resp
 
     def get_poi_categories(
@@ -153,6 +204,23 @@ class Simulator:
         center: Optional[Union[tuple[float, float], Point]] = None,
         radius: Optional[float] = None,
     ) -> list[str]:
+        """
+        Retrieve unique categories of Points of Interest (POIs) around a central point.
+
+        - **Args**:
+            - `center` (`Optional[Union[Tuple[float, float], Point]]`): The central point as a tuple or Point object.
+              Defaults to (0, 0) if not provided.
+            - `radius` (`Optional[float]`): The search radius in meters. If not provided, all POIs are considered.
+
+        - **Returns**:
+            - `List[str]`: A list of unique POI category names.
+        """
+        start_time = time.time()
+        log = {
+            "req": "get_poi_categories",
+            "start_time": start_time,
+            "consumption": 0
+        }
         categories: list[str] = []
         if center is None:
             center = (0, 0)
@@ -164,23 +232,31 @@ class Simulator:
         for poi in _pois:
             catg = poi["category"]
             categories.append(catg.split("|")[-1])
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
         return list(set(categories))
 
     async def get_time(
         self, format_time: bool = False, format: str = "%H:%M:%S"
     ) -> Union[int, str]:
         """
-        获取模拟器当前时间 Get current time of simulator
-        默认返回以00:00:00为始的, 以s为单位的时间(int)
-        支持格式化时间
+        Get the current time of the simulator.
 
-        Args:
-        - format_time (bool): 是否格式化 format or not
-        - format (str): 格式化模板，默认为"Hour:Minute:Second" the formation
+        By default, returns the number of seconds since midnight. Supports formatted output.
 
-        Returns:
-        - time Union[int, str]: 时间 time in second(int) or formatted time(str)
+        - **Args**:
+            - `format_time` (`bool`): Whether to return the time in a formatted string. Defaults to `False`.
+            - `format` (`str`): The format string for formatting the time. Defaults to "%H:%M:%S".
+
+        - **Returns**:
+            - `Union[int, str]`: The current simulation time either as an integer representing seconds since midnight or as a formatted string.
         """
+        start_time = time.time()
+        log = {
+            "req": "get_time",
+            "start_time": start_time,
+            "consumption": 0
+        }
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
         self.time = now["t"]
@@ -189,44 +265,133 @@ class Simulator:
             start_of_day = datetime.combine(current_date, datetime.min.time())
             current_time = start_of_day + timedelta(seconds=now["t"])
             formatted_time = current_time.strftime(format)
+            log["consumption"] = time.time() - start_time
+            self._log_list.append(log)
             return formatted_time
         else:
+            log["consumption"] = time.time() - start_time
+            self._log_list.append(log)
             return int(now["t"])
 
     async def pause(self):
+        """
+        Pause the simulation.
+
+        This method sends a request to the simulator's pause service to pause the simulation.
+        """
+        start_time = time.time()
+        log = {
+            "req": "pause",
+            "start_time": start_time,
+            "consumption": 0
+        }
         await self._client.pause_service.pause()
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
 
     async def resume(self):
+        """
+        Resume the simulation.
+
+        This method sends a request to the simulator's pause service to resume the simulation.
+        """
+        start_time = time.time()
+        log = {
+            "req": "resume",
+            "start_time": start_time,
+            "consumption": 0
+        }
         await self._client.pause_service.resume()
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
 
     async def get_simulator_day(self) -> int:
         """
-        获取模拟器到第几日
+        Get the current day of the simulation.
+
+        - **Returns**:
+            - `int`: The day number since the start of the simulation.
         """
+        start_time = time.time()
+        log = {
+            "req": "get_simulator_day",
+            "start_time": start_time,
+            "consumption": 0
+        }
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
         day = now["day"]
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
         return day
 
     async def get_simulator_second_from_start_of_day(self) -> int:
         """
-        获取模拟器从00:00:00到当前的秒数
+        Get the number of seconds elapsed from the start of the current day in the simulation.
+
+        - **Returns**:
+            - `int`: The number of seconds from 00:00:00 of the current day.
         """
+        start_time = time.time()
+        log = {
+            "req": "get_simulator_second_from_start_of_day",
+            "start_time": start_time,
+            "consumption": 0
+        }
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
         return now["t"] % 86400
 
     async def get_person(self, person_id: int) -> dict:
-        return await self._client.person_service.GetPerson(
+        """
+        Retrieve information about a specific person by ID.
+
+        - **Args**:
+            - `person_id` (`int`): The ID of the person to retrieve information for.
+
+        - **Returns**:
+            - `Dict`: Information about the specified person.
+        """
+        start_time = time.time()
+        log = {
+            "req": "get_person",
+            "start_time": start_time,
+            "consumption": 0
+        }
+        person = await self._client.person_service.GetPerson(
             req={"person_id": person_id}
         )  # type:ignore
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
+        return person
 
     async def add_person(self, person: Any) -> dict:
+        """
+        Add a new person to the simulation.
+
+        - **Args**:
+            - `person` (`Any`): The person object to add. If it's an instance of `person_pb2.Person`,
+              it will be wrapped in an `AddPersonRequest`. Otherwise, `person` is expected to already be a valid request object.
+
+        - **Returns**:
+            - `Dict`: Response from adding the person.
+        """
+        start_time = time.time()
+        log = {
+            "req": "add_person",
+            "start_time": start_time,
+            "consumption": 0
+        }
         if isinstance(person, person_pb2.Person):
             req = person_service.AddPersonRequest(person=person)
         else:
             req = person
-        return await self._client.person_service.AddPerson(req)  # type:ignore
+        person_id = await self._client.person_service.AddPerson(req)  # type:ignore
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
+        return person_id
 
     async def set_aoi_schedules(
         self,
@@ -237,6 +402,24 @@ class Simulator:
         departure_times: Optional[list[float]] = None,
         modes: Optional[list[TripMode]] = None,
     ):
+        """
+        Set schedules for a person to visit Areas of Interest (AOIs).
+
+        - **Args**:
+            - `person_id` (`int`): The ID of the person whose schedule is being set.
+            - `target_positions` (`Union[List[Union[int, Tuple[int, int]]], Union[int, Tuple[int, int]]]`):
+              A list of AOI or POI IDs or tuples of (AOI ID, POI ID) that the person will visit.
+            - `departure_times` (`Optional[List[float]]`): Departure times for each trip in the schedule.
+              If not provided, current time will be used for all trips.
+            - `modes` (`Optional[List[TripMode]]`): Travel modes for each trip.
+              Defaults to `TRIP_MODE_DRIVE_ONLY` if not specified.
+        """
+        start_time = time.time()
+        log = {
+            "req": "set_aoi_schedules",
+            "start_time": start_time,
+            "consumption": 0
+        }
         cur_time = float(await self.get_time())
         if not isinstance(target_positions, list):
             target_positions = [target_positions]
@@ -286,6 +469,8 @@ class Simulator:
             )
         req = {"person_id": person_id, "schedules": _schedules}
         await self._client.person_service.SetSchedule(req)
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
 
     async def reset_person_position(
         self,
@@ -295,6 +480,22 @@ class Simulator:
         lane_id: Optional[int] = None,
         s: Optional[float] = None,
     ):
+        """
+        Reset the position of a person within the simulation.
+
+        - **Args**:
+            - `person_id` (`int`): The ID of the person whose position is being reset.
+            - `aoi_id` (`Optional[int]`): The ID of the Area of Interest (AOI) where the person should be placed.
+            - `poi_id` (`Optional[int]`): The ID of the Point of Interest (POI) within the AOI.
+            - `lane_id` (`Optional[int]`): The ID of the lane on which the person should be placed.
+            - `s` (`Optional[float]`): The longitudinal position along the lane.
+        """
+        start_time = time.time()
+        log = {
+            "req": "reset_person_position",
+            "start_time": start_time,
+            "consumption": 0
+        }
         reset_position = {}
         if aoi_id is not None:
             reset_position["aoi_position"] = {"aoi_id": aoi_id}
@@ -323,6 +524,8 @@ class Simulator:
             logger.debug(
                 f"Neither aoi or lane pos provided for person {person_id} position reset!!"
             )
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
 
     def get_around_poi(
         self,
@@ -330,9 +533,26 @@ class Simulator:
         radius: float,
         poi_type: Union[str, list[str]],
     ) -> list[dict]:
+        """
+        Get Points of Interest (POIs) around a central point based on type.
+
+        - **Args**:
+            - `center` (`Union[Tuple[float, float], Point]`): The central point as a tuple or Point object.
+            - `radius` (`float`): The search radius in meters.
+            - `poi_type` (`Union[str, List[str]]`): The category or categories of POIs to filter by.
+
+        - **Returns**:
+            - `List[Dict]`: A list of dictionaries containing information about the POIs found.
+        """
+        start_time = time.time()
+        log = {
+            "req": "get_around_poi",
+            "start_time": start_time,
+            "consumption": 0
+        }
         if isinstance(poi_type, str):
             poi_type = [poi_type]
-        transformed_poi_type = []
+        transformed_poi_type: list[str] = []
         for t in poi_type:
             if t not in self.poi_cate:
                 transformed_poi_type.append(t)
@@ -352,4 +572,6 @@ class Simulator:
             if catg.split("|")[-1] not in poi_type_set:
                 continue
             pois.append(poi)
-        return pois
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
+        return pois

pycityagent/environment/utils/__init__.py

@@ -4,11 +4,11 @@ utilities
 """
 
 from .geojson import wrap_feature_collection
-from .base64 import encode_to_base64
 from .port import find_free_port
+from .base64 import encode_to_base64
 
 __all__ = [
     "wrap_feature_collection",
     "find_free_port",
-    "find_free_port",
+    "encode_to_base64",
 ]

pycityagent/environment/utils/geojson.py

@@ -6,11 +6,11 @@ def wrap_feature_collection(features: list[dict], name: str):
     将 GeoJSON Feature 集合包装为 FeatureCollection
     Wrap GeoJSON Feature collection as FeatureCollection
 
-    Args:
+    - **Args**:
     - features: GeoJSON Feature 集合。GeoJSON Feature collection.
     - name: FeatureCollection 名称。FeatureCollection name.
 
-    Returns:
+    - **Returns**:
     - dict: GeoJSON FeatureCollection
     """
     return {

pycityagent/environment/utils/grpc.py

@@ -8,11 +8,11 @@ def create_channel(server_address: str, secure: bool = False) -> grpc.Channel:
     创建一个grpc的channel
     Create a grpc channel
 
-    Args:
+    - **Args**:
     - server_address (str): 服务器地址。server address.
     - secure (bool, optional): 是否使用安全连接. Defaults to False. Whether to use a secure connection. Defaults to False.
 
-    Returns:
+    - **Returns**:
     - grpc.Channel: grpc的channel。grpc channel.
     """
     if server_address.startswith("http://"):
@@ -35,11 +35,11 @@ def create_aio_channel(server_address: str, secure: bool = False) -> grpc.aio.Ch
     创建一个grpc的异步channel
     Create a grpc asynchronous channel
 
-    Args:
+    - **Args**:
     - server_address (str): 服务器地址。server address.
     - secure (bool, optional): 是否使用安全连接. Defaults to False. Whether to use a secure connection. Defaults to False.
 
-    Returns:
+    - **Returns**:
     - grpc.aio.Channel: grpc的异步channel。grpc asynchronous channel.
     """
     if server_address.startswith("http://"):

pycityagent/environment/utils/map_utils.py

@@ -13,12 +13,12 @@ def point_on_line_given_distance(start_node, end_node, distance):
     Given two points (start_point and end_point) defining a line, and a distance s to travel along the line,
     return the coordinates of the point reached after traveling s units along the line, starting from start_point.
 
-    Args:
+    - **Args**:
         start_point (tuple): tuple of (x, y) representing the starting point on the line.
         end_point (tuple): tuple of (x, y) representing the ending point on the line.
         distance (float): Distance to travel along the line, starting from start_point.
 
-    Returns:
+    - **Returns**:
         tuple: tuple of (x, y) representing the new point reached after traveling s units along the line.
     """