migrate routing command to use the models and store data there

crawler/models/listing.py

@@ -1,9 +1,13 @@
-from dataclasses import dataclass
-from datetime import datetime
+from __future__ import annotations
+from dataclasses import asdict, dataclass
+import dataclasses
+from datetime import datetime, timedelta
 import enum
+import json
 from pathlib import Path
 from typing import Any, Dict, List
-from sqlmodel import JSON, Column, Enum, SQLModel, Field, String, TypeDecorator
+from rec import routing
+from sqlmodel import JSON, SQLModel, Field, String
 
 
 @dataclass
@@ -13,6 +17,28 @@ class PriceHistoryItem:
     price: float
 
 
+@dataclass(frozen=True)
+class Route:
+    legs: list[RouteLegStep]
+    distance_meters: int
+    duration_s: int
+
+    @property
+    def duration(self) -> timedelta:
+        return timedelta(seconds=self.duration_s)
+
+
+@dataclass(frozen=True)
+class RouteLegStep:
+    distance_meters: int
+    duration_s: int
+    travel_mode: routing.TravelMode
+
+    @property
+    def duration(self) -> timedelta:
+        return timedelta(seconds=self.duration_s)
+
+
 class ListingSite(enum.StrEnum):
     RIGHTMOVE = "rightmove"
     # ZOOPLA = "zoopla"
@@ -38,6 +64,69 @@ class Listing(SQLModel, table=False):
     additional_info: Dict[str, Any] = Field(
         default_factory=dict, sa_type=JSON, nullable=False
     )
+    routing_info_json: str = Field(
+        sa_type=String, nullable=True, default=None
+    )  # Store as JSON string for simplicity
+
+    @property
+    def is_removed(self) -> bool:
+        return not self.additional_info["property"]["visible"]
+
+    @property
+    def routing_info(self) -> dict[DestinationMode, List[Route]]:
+        """
+        Returns a list of DestinationMode objects from the routing_info_str.
+        """
+        if not self.routing_info_json:
+            return {}
+
+        # TODO: move to a separate serializer class
+        json_data = json.loads(self.routing_info_json)
+        destimation_routes = {}
+        for destination_mode_str, routes_json in json_data.items():
+            destination_mode = DestinationMode(
+                destination_address=json.loads(destination_mode_str)[
+                    "destination_address"
+                ],
+                travel_mode=routing.TravelMode(
+                    json.loads(destination_mode_str)["travel_mode"]
+                ),
+            )
+            parsed_route = json.loads(routes_json[0])
+            routes = [
+                Route(
+                    legs=[
+                        RouteLegStep(
+                            distance_meters=step["distance_meters"],
+                            duration_s=step["duration_s"],
+                            travel_mode=routing.TravelMode(step["travel_mode"]),
+                        )
+                        for step in parsed_route["legs"]
+                    ],
+                    distance_meters=parsed_route["distance_meters"],
+                    duration_s=int(parsed_route["duration_s"]),
+                )
+            ]
+            destimation_routes[destination_mode] = routes
+        return destimation_routes
+
+    def serialize_routing_info(
+        self, routing_info: dict[DestinationMode, list[Route]]
+    ) -> str:
+        """
+        Serializes the routing_info to a JSON string.
+        """
+        # TODO: move to a separate serializer class
+        # for destination_mode, routes in routing_info.items():
+        serialized = json.dumps(
+            {
+                json.dumps(dataclasses.asdict(destination_mode)): [
+                    json.dumps(dataclasses.asdict(route)) for route in routes
+                ]
+                for destination_mode, routes in routing_info.items()
+            }
+        )
+        return serialized
 
 
 class FurnishType(enum.StrEnum):
@@ -57,3 +146,39 @@ class BuyListing(Listing, table=True):
     lease_left: int | None = Field(
         default=None, nullable=True
     )  # in years, e.g., 90, 80, etc.
+
+
+@dataclass(frozen=True)
+class DestinationMode:
+    destination_address: str
+    travel_mode: routing.TravelMode
+
+    def __hash__(self) -> int:
+        return hash((self.destination_address, self.travel_mode))
+
+    # def to_dict(self) -> dict[str, str | routing.TravelMode]:
+    #     return {
+    #         "destination_address": self.destination_address,
+    #         "travel_mode": self.travel_mode.value,
+    #     }
+
+    # @classmethod
+    # def from_dict(cls, data: dict):
+    #     return cls(
+    #         destination_address=data["destination_address"],
+    #         travel_mode=routing.TravelMode(data["travel_mode"]),
+    #     )
+
+    # def __json__(self) -> dict[str, str | routing.TravelMode]:
+    #     return {
+    #         "destination_address": self.destination_address,
+    #         "travel_mode": self.travel_mode.value,
+    #     }
+
+    def __getstate__(self):
+        # This allows serializers to pick up a dict representation
+        return asdict(self)
+
+    def __iter__(self):
+        # Makes it behave like a dict when expected
+        return iter(asdict(self).items())