refactor Timeline class to pydantic model

Author: tanwinn

app/models.py

@@ -1,7 +1,7 @@
 """app.models.py"""
 from typing import Dict, List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
 
 
 class Latest(BaseModel):
@@ -27,8 +27,25 @@ class Timeline(BaseModel):
     Timeline model.
     """
 
-    latest: int
-    timeline: Dict[str, int] = {}
+    timeline: Dict[str, int] = dict()
+
+    @validator("timeline")
+    @classmethod
+    def sort_timeline(cls, value):
+        """Sort the timeline history before inserting into the model"""
+        return dict(sorted(value.items()))
+
+    @property
+    def latest(self):
+        """Get latest available history value."""
+        return list(self.timeline.values())[-1] if self.timeline else 0
+
+    def serialize(self):
+        """
+        Serialize the model into dict
+        TODO: override dict() instead of using serialize
+        """
+        return {**self.dict(), "latest": self.latest}
 
 
 class Timelines(BaseModel):

app/services/location/csbs.py

@@ -97,6 +97,6 @@ async def get_locations():
             await load_cache(data_id, locations)
         except TypeError as type_err:
             LOGGER.error(type_err)
-            
+
     # Return the locations.
     return locations

app/services/location/jhu.py

@@ -11,7 +11,7 @@
 from ...caches import check_cache, load_cache
 from ...coordinates import Coordinates
 from ...location import TimelinedLocation
-from ...timeline import Timeline
+from ...models import Timeline
 from ...utils import countries
 from ...utils import date as date_util
 from ...utils import httputils
@@ -175,18 +175,18 @@ async def get_locations():
                 # Timelines (parse dates as ISO).
                 {
                     "confirmed": Timeline(
-                        {
+                        timeline={
                             datetime.strptime(date, "%m/%d/%y").isoformat() + "Z": amount
                             for date, amount in timelines["confirmed"].items()
                         }
                     ),
                     "deaths": Timeline(
-                        {
+                        timeline={
                             datetime.strptime(date, "%m/%d/%y").isoformat() + "Z": amount
                             for date, amount in timelines["deaths"].items()
                         }
                     ),
-                    "recovered": Timeline({}),
+                    "recovered": Timeline(),
                 },
             )
         )

app/services/location/nyt.py

@@ -9,7 +9,7 @@
 from ...caches import check_cache, load_cache
 from ...coordinates import Coordinates
 from ...location.nyt import NYTLocation
-from ...timeline import Timeline
+from ...models import Timeline
 from ...utils import httputils
 from . import LocationService
 
@@ -119,18 +119,18 @@ async def get_locations():
                     last_updated=datetime.utcnow().isoformat() + "Z",  # since last request
                     timelines={
                         "confirmed": Timeline(
-                            {
+                            timeline={
                                 datetime.strptime(date, "%Y-%m-%d").isoformat() + "Z": amount
                                 for date, amount in confirmed_history.items()
                             }
                         ),
                         "deaths": Timeline(
-                            {
+                            timeline={
                                 datetime.strptime(date, "%Y-%m-%d").isoformat() + "Z": amount
                                 for date, amount in deaths_history.items()
                             }
                         ),
-                        "recovered": Timeline({}),
+                        "recovered": Timeline(),
                     },
                 )
             )

app/timeline.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from app import coordinates, location, timeline
+from app import coordinates, location, models
 
 
 def mocked_timeline(*args, **kwargs):
@@ -22,7 +22,7 @@ def __init__(self, latest):
         (2, "Cruise Ship", "XX", "", 15, 100, 1000, 1111, 22222),
     ],
 )
-@mock.patch("app.timeline.Timeline", side_effect=mocked_timeline)
+@mock.patch("app.models.Timeline", side_effect=mocked_timeline)
 def test_location_class(
     mocked_timeline,
     test_id,
@@ -39,9 +39,9 @@ def test_location_class(
     coords = coordinates.Coordinates(latitude=latitude, longitude=longitude)
 
     # Timelines
-    confirmed = timeline.Timeline(confirmed_latest)
-    deaths = timeline.Timeline(deaths_latest)
-    recovered = timeline.Timeline(recovered_latest)
+    confirmed = models.Timeline(confirmed_latest)
+    deaths = models.Timeline(deaths_latest)
+    recovered = models.Timeline(recovered_latest)
 
     # Date now.
     now = datetime.utcnow().isoformat() + "Z"

tests/test_location.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from app import timeline
+from app import models
 
 
 def test_timeline_class():
@@ -15,7 +15,7 @@ def test_timeline_class():
         "1/23/20": 3,
     }
 
-    history_data = timeline.Timeline(history=timeseries)
+    history_data = models.Timeline(timeline=timeseries)
 
     # validate last value
     assert history_data.latest == 7