feat: pydantic

Author: Strvm

README.md

@@ -7,6 +7,8 @@ RatpAPI is a Python wrapper for interacting with the RATP (Régie Autonome des T
 - Fetch global traffic information on the RATP network.
 - Retrieve traffic information for specific lines using the `LineID` enum for easy reference.
 - Get affluence data for journeys on particular lines.
+- Utilizes Pydantic models for data validation and parsing.
+
 
 ## Installation
 To use RatpAPI in your project, you can install it via pip:
@@ -15,7 +17,7 @@ pip install ratp-api-python
 ```
 
 ## Usage
-
+a
 ### Setting Up
 First, import the `RatpAPI` class and initialize it with your API key:
 
@@ -44,8 +46,14 @@ print(global_traffic)
 ### Fetching Line-Specific Traffic
 To get traffic information for a specific line using the `LineID` enum:
 ```python
-line_traffic = api.get_line_traffic(LineID.RER_A)
-print(line_traffic)
+line_data = api.get_line_traffic(line_id=LineID.METRO_14)
+for situation in line_data.situations:
+    print(situation)
+# Output:
+# isActive=True isPlanned=True criticity='HIGH' messages=["Jusqu'au 04/02, le week-end, trafic interrompu sur l'ensemble de la ligne en raison de travaux. Bus de remplacement. Plus d'informations sur la page dédiée."]
+# isActive=False isPlanned=True criticity='HIGH' messages=["Du 01/02 au 08/02, du lundi au jeudi à partir de 22h, trafic interrompu sur l'ensemble de la ligne en raison de travaux. Bus de remplacement. Plus d'informations sur la page dédiée."]
+# isActive=False isPlanned=True criticity='HIGH' messages=["Jusqu'au 31/01, du lundi au jeudi à partir de 22h, trafic interrompu sur l'ensemble de la ligne en raison de travaux. Bus de remplacement. Plus d'informations sur la page dédiée."]
+
 ```
 
 ### Fetching Affluence for Journeys

src/ratp_api/main.py

@@ -5,6 +5,7 @@
 
 import requests
 from affluence.line_mappings import line_end_points
+from ratp_api.models.line import LineData
 
 logging.basicConfig(level=logging.INFO)
 
@@ -122,15 +123,17 @@ def get_global_traffic(self) -> dict:
         """
         return self.make_request("GET", self.GLOBAL_TRAFFIC_URL)
 
-    def get_line_traffic(self, line_id: LineID) -> dict:
+    def get_line_traffic(self, line_id: LineID) -> LineData:
         """
         Retrieve traffic information for a specific line.
 
         :param line_id: The ID of the line for which traffic information is requested.
         :return: A dictionary representing the traffic situation for the specified line.
         """
         url = self.LINE_TRAFFIC_URL.format(line_id)
-        return self.make_request("GET", url)
+        line_data_raw = self.make_request("GET", url)
+        line_data = LineData(**line_data_raw)
+        return line_data
 
     def get_affluence(self, start, end, line_id):
         """
@@ -210,7 +213,39 @@ def get_all_lines_affluence(self):
 
         return affluences
 
+    def get_line_affluence(self, line_id: LineID):
+        """
+        Retrieve affluence information for a specific line.
+
+        :param line_id: The ID of the line for which affluence information is requested.
+        :return: A dictionary containing affluence information, or an empty dict if not available.
+        """
+        line_id = str(line_id.name).split("_")[1]
+        endpoints = line_end_points.get(line_id)
+        if not endpoints:
+            logging.info(f"No endpoints found for {line_id}")
+            return {}
+        try:
+            comfort = self.get_affluence(
+                start=endpoints["start"]["coords"],
+                end=endpoints["end"]["coords"],
+                line_id=line_id,
+            )
+            if "level" in comfort:
+                comfort["lineDisplayCode"] = line_id
+                return comfort
+            else:
+                logging.info(f"No affluence level found for {line_id}")
+        except Exception as e:
+            logging.info(f"Error fetching affluence for {line_id}: {e}")
+        return {}
+
 
 if __name__ == "__main__":
     api = RatpAPI(api_key="e2rDkJzd2c1dPaFh7e0pJ9H7NjeqTQHg6ql31LmZ")
-    print(api.get_line_traffic(line_id=LineID.METRO_1))
+    line_data = api.get_line_traffic(line_id=LineID.METRO_14)
+    print(line_data)
+    for situation in line_data.situations:
+        print(situation)
+    # print(api.get_line_affluence(line_id=LineID.METRO_9))
+    # print(api.get_all_lines_affluence())

src/ratp_api/models/__init__.py

@@ -0,0 +1,38 @@
+from pydantic import BaseModel, HttpUrl
+
+
+class Icon(BaseModel):
+    svg: HttpUrl
+    png: HttpUrl
+    name: str
+
+
+class Assets(BaseModel):
+    icon: Icon
+
+
+class Line(BaseModel):
+    id: str
+    name: str
+    displayCode: str
+    businessMode: str
+    assets: Assets
+
+
+class SituationMessage(BaseModel):
+    isActive: bool
+    isPlanned: bool
+    criticity: str
+    messages: list[str]
+
+
+class TwitterAccountInfos(BaseModel):
+    twitterId: str
+    twitterUrl: HttpUrl
+    isThreadActive: bool
+
+
+class LineData(BaseModel):
+    line: Line
+    situations: list[SituationMessage]
+    twitterAccountInfos: TwitterAccountInfos