Use pydantic for JSON parsing

nextmv/base_model.py

@@ -0,0 +1,20 @@
+"""JSON class for data wrangling JSON objects."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class BaseModel(BaseModel):
+    """Base class for data wrangling tasks with JSON."""
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Instantiates the class from a dict."""
+
+        return cls(**data)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Converts the class to a dict."""
+
+        return self.model_dump(mode="json", exclude_none=True)

nextmv/cloud/application.py

@@ -4,12 +4,11 @@
 from datetime import datetime
 from typing import Any
 
+from nextmv.base_model import BaseModel
 from nextmv.cloud.client import Client
-from nextmv.json_class import JSONClass
 
 
-@dataclass
-class Metadata(JSONClass):
+class Metadata(BaseModel):
     """Metadata of a run, whether it was successful or not."""
 
     status: str
@@ -32,8 +31,7 @@ class Metadata(JSONClass):
     """ID of the version of the application where the run was submitted to."""
 
 
-@dataclass
-class RunResult(JSONClass):
+class RunResult(BaseModel):
     """Result of a run, wheter it was successful or not."""
 
     id: str
@@ -89,6 +87,9 @@ def new_run(
             description: Description of the run.
             upload_id: ID to use when running a large input.
             options: Options to use for the run.
+
+        Returns:
+            ID of the submitted run.
         """
 
         payload = {
@@ -124,6 +125,9 @@ def run_result(
 
         Args:
             run_id: ID of the run.
+
+        Returns:
+            Result of the run.
         """
 
         response = self.client.get(

nextmv/cloud/client.py

@@ -44,7 +44,17 @@ def post(
         payload: dict[str, Any],
         query_params: dict[str, Any] | None = None,
     ) -> requests.Response:
-        """Send a POST request to the Nextmv Cloud API."""
+        """
+        Send a POST request to the Nextmv Cloud API.
+
+        Args:
+            endpoint: Endpoint to send the request to.
+            payload: Payload to send with the request.
+            query_params: Query parameters to send with the request.
+
+        Returns:
+            Response from the Nextmv Cloud API.
+        """
 
         return requests.post(
             url=f"{self.url}/{endpoint}",
@@ -58,7 +68,16 @@ def get(
         endpoint: str,
         query_params: dict[str, Any] | None = None,
     ) -> requests.Response:
-        """Send a GET request to the Nextmv Cloud API."""
+        """
+        Send a GET request to the Nextmv Cloud API.
+
+        Args:
+            endpoint: Endpoint to send the request to.
+            query_params: Query parameters to send with the request.
+
+        Returns:
+            Response from the Nextmv Cloud API.
+        """
 
         return requests.get(
             url=f"{self.url}/{endpoint}",

nextmv/nextroute/schema/input.py

@@ -1,15 +1,13 @@
 """Defines the input class"""
 
-from dataclasses import dataclass
 from typing import Any
 
-from nextmv.json_class import JSONClass
+from nextmv.base_model import BaseModel
 from nextmv.nextroute.schema.stop import AlternateStop, Stop, StopDefaults
 from nextmv.nextroute.schema.vehicle import Vehicle, VehicleDefaults
 
 
-@dataclass
-class Defaults(JSONClass):
+class Defaults(BaseModel):
     """Default values for vehicles and stops."""
 
     stops: StopDefaults | None = None
@@ -18,8 +16,7 @@ class Defaults(JSONClass):
     """Default values for vehicles."""
 
 
-@dataclass
-class DurationGroup(JSONClass):
+class DurationGroup(BaseModel):
     """Represents a group of stops that get additional duration whenever a stop
     of the group is approached for the first time."""
 
@@ -29,8 +26,7 @@ class DurationGroup(JSONClass):
     """Stop IDs contained in the group."""
 
 
-@dataclass
-class Input(JSONClass):
+class Input(BaseModel):
     """Input schema for Nextroute."""
 
     stops: list[Stop]

nextmv/nextroute/schema/location.py

@@ -1,13 +1,10 @@
 """Defines the location class."""
 
 
-from dataclasses import dataclass
+from nextmv.base_model import BaseModel
 
-from nextmv.json_class import JSONClass
 
-
-@dataclass
-class Location(JSONClass):
+class Location(BaseModel):
     """Location represents a geographical location."""
 
     lon: float

nextmv/nextroute/schema/stop.py

@@ -1,16 +1,14 @@
 """Defines the stop class."""
 
 
-from dataclasses import dataclass
 from datetime import datetime
 from typing import Any
 
-from nextmv.json_class import JSONClass
+from nextmv.base_model import BaseModel
 from nextmv.nextroute.schema.location import Location
 
 
-@dataclass
-class StopDefaults(JSONClass):
+class StopDefaults(BaseModel):
     """Default values for a stop."""
 
     compatibility_attributes: list[str] | None = None
@@ -33,7 +31,6 @@ class StopDefaults(JSONClass):
     """Penalty for not planning a stop."""
 
 
-@dataclass(kw_only=True)
 class Stop(StopDefaults):
     """Stop is a location that must be visited by a vehicle in a Vehicle
     Routing Problem (VRP.)"""
@@ -53,7 +50,6 @@ class Stop(StopDefaults):
     """Stops that must be visited before this one on the same route."""
 
 
-@dataclass(kw_only=True)
 class AlternateStop(StopDefaults):
     """An alternate stop can be serviced instead of another stop."""
 

nextmv/nextroute/schema/vehicle.py

@@ -1,16 +1,14 @@
 """Defines the vehicle class."""
 
 
-from dataclasses import dataclass
 from datetime import datetime
 from typing import Any
 
-from nextmv.json_class import JSONClass
+from nextmv.base_model import BaseModel
 from nextmv.nextroute.schema.location import Location
 
 
-@dataclass
-class InitialStop(JSONClass):
+class InitialStop(BaseModel):
     """Represents a stop that is already planned on a vehicle."""
 
     id: str
@@ -20,8 +18,7 @@ class InitialStop(JSONClass):
     """Whether the stop is fixed on the route."""
 
 
-@dataclass
-class VehicleDefaults(JSONClass):
+class VehicleDefaults(BaseModel):
     """Default values for vehicles."""
 
     activation_penalty: int | None = None
@@ -58,7 +55,6 @@ class VehicleDefaults(JSONClass):
     """Time when the vehicle starts its route."""
 
 
-@dataclass(kw_only=True)
 class Vehicle(VehicleDefaults):
     """A vehicle services stops in a Vehicle Routing Problem (VRP)."""
 

requirements.txt

@@ -1,3 +1,4 @@
+pydantic==2.5.2
 requests==2.31.0
 ruff==0.1.7
 setuptools==69.0.2

tests/nextroute/schema/test_input.py

@@ -2,40 +2,56 @@
 import os
 import unittest
 
-from nextmv.nextroute.schema import Input, Stop
+from nextmv.nextroute.schema import Input, Stop, Vehicle
 
 
 class TestInput(unittest.TestCase):
+    filepath = os.path.join(os.path.dirname(__file__), "input.json")
+
     def test_from_json(self):
-        filepath = os.path.join(os.path.dirname(__file__), "input.json")
-        input = Input.from_file(filepath)
-        parsed = input.to_dict()
-        with open(filepath) as f:
-            expected = json.load(f)
+        with open(self.filepath) as f:
+            json_data = json.load(f)
+
+        nextroute_input = Input.from_dict(json_data)
+        parsed = nextroute_input.to_dict()
+
+        for s, stop in enumerate(parsed["stops"]):
+            original_stop = json_data["stops"][s]
+            self.assertEqual(
+                stop,
+                original_stop,
+                f"stop: parsed({stop}) and original ({original_stop}) should be equal",
+            )
+
+        for v, vehicle in enumerate(parsed["vehicles"]):
+            original_vehicle = json_data["vehicles"][v]
+            self.assertEqual(
+                vehicle,
+                original_vehicle,
+                f"vehicle: parsed ({vehicle}) and original ({original_vehicle}) should be equal",
+            )
 
-        self.maxDiff = None
         self.assertEqual(
-            parsed,
-            expected,
-            "Parsing the JSON into the class and back should yield the same result.",
+            parsed["defaults"],
+            json_data["defaults"],
+            f"defaults: parsed ({parsed['defaults']}) and original ({json_data['defaults']}) should be equal",
         )
 
     def test_from_dict(self):
-        filepath = os.path.join(os.path.dirname(__file__), "input.json")
-        with open(filepath) as f:
-            expected = json.load(f)
+        with open(self.filepath) as f:
+            json_data = json.load(f)
 
-        input = Input.from_dict(expected)
-        stops = input.stops
+        nextroute_input = Input.from_dict(json_data)
+        stops = nextroute_input.stops
         for stop in stops:
             self.assertTrue(
                 isinstance(stop, Stop),
-                f"Stop {stop} should be of type Stop.",
+                f"Stop ({stop}) should be of type Stop.",
             )
 
-        vehicles = input.vehicles
+        vehicles = nextroute_input.vehicles
         for vehicle in vehicles:
             self.assertTrue(
-                isinstance(vehicle, Stop),
-                f"Vehicle {vehicle} should be of type Vehicle.",
+                isinstance(vehicle, Vehicle),
+                f"Vehicle ({vehicle}) should be of type Vehicle.",
             )