Apply Structural Design Pattern

Author: dengrf96

app/location/__init__.py

@@ -11,7 +11,7 @@ class Location:  # pylint: disable=too-many-instance-attributes
     """
 
     def __init__(
-        self, id, country, province, coordinates, last_updated, confirmed, deaths, recovered,
+            self, id, country, province, coordinates, last_updated, confirmed, deaths, recovered, timelines
     ):  # pylint: disable=too-many-arguments
         # General info.
         self.id = id
@@ -26,6 +26,7 @@ def __init__(
         self.confirmed = confirmed
         self.deaths = deaths
         self.recovered = recovered
+        self.timelines = timelines
 
     @property
     def country_code(self):
@@ -47,14 +48,14 @@ def country_population(self):
         """
         return country_population(self.country_code)
 
-    def serialize(self):
+    def serialize(self, timelines=False):
         """
         Serializes the location into a dict.
 
         :returns: The serialized location.
         :rtype: dict
         """
-        return {
+        serialized = {
             # General info.
             "id": self.id,
             "country": self.country,
@@ -72,43 +73,6 @@ def serialize(self):
                 "recovered": self.recovered,
             },
         }
-
-
-class TimelinedLocation(Location):
-    """
-    A location with timelines.
-    """
-
-    # pylint: disable=too-many-arguments
-    def __init__(self, id, country, province, coordinates, last_updated, timelines):
-        super().__init__(
-            # General info.
-            id,
-            country,
-            province,
-            coordinates,
-            last_updated,
-            # Statistics (retrieve latest from timelines).
-            confirmed=timelines.get("confirmed").latest or 0,
-            deaths=timelines.get("deaths").latest or 0,
-            recovered=timelines.get("recovered").latest or 0,
-        )
-
-        # Set timelines.
-        self.timelines = timelines
-
-    # pylint: disable=arguments-differ
-    def serialize(self, timelines=False):
-        """
-        Serializes the location into a dict.
-
-        :param timelines: Whether to include the timelines.
-        :returns: The serialized location.
-        :rtype: dict
-        """
-        serialized = super().serialize()
-
-        # Whether to include the timelines or not.
         if timelines:
             serialized.update(
                 {
@@ -119,6 +83,4 @@ def serialize(self, timelines=False):
                     }
                 }
             )
-
-        # Return the serialized location.
         return serialized

app/location/csbs.py

@@ -20,6 +20,7 @@ def __init__(self, id, state, county, coordinates, last_updated, confirmed, deat
             confirmed=confirmed,
             deaths=deaths,
             recovered=0,
+            timelines=None
         )
 
         self.state = state
@@ -32,7 +33,7 @@ def serialize(self, timelines=False):  # pylint: disable=arguments-differ,unused
         :returns: The serialized location.
         :rtype: dict
         """
-        serialized = super().serialize()
+        serialized = super().serialize(timelines)
 
         # Update with new fields.
         serialized.update(

app/location/jhu.py

@@ -0,0 +1,23 @@
+"""app.locations.nyt.py"""
+from app.location import Location
+
+
+class JHULocation(Location):
+    """
+    A JHU (county) Location.
+    """
+
+    # pylint: disable=too-many-arguments,redefined-builtin
+    def __init__(self, id, country, state, coordinates, last_updated, timelines):
+        super().__init__(
+            # General info.
+            id,
+            country,
+            state,
+            coordinates,
+            last_updated,
+            confirmed=timelines.get("confirmed").latest or 0,
+            deaths=timelines.get("deaths").latest or 0,
+            recovered=timelines.get("recovered").latest or 0,
+            timelines=timelines
+        )

app/location/nyt.py

@@ -1,15 +1,26 @@
 """app.locations.nyt.py"""
-from . import TimelinedLocation
+from app.location import Location
 
 
-class NYTLocation(TimelinedLocation):
+class NYTLocation(Location):
     """
-    A NYT (county) Timelinedlocation.
+    A NYT (county) Location.
     """
 
     # pylint: disable=too-many-arguments,redefined-builtin
     def __init__(self, id, state, county, coordinates, last_updated, timelines):
-        super().__init__(id, "US", state, coordinates, last_updated, timelines)
+        super().__init__(
+            # General info.
+            id,
+            "US",
+            state,
+            coordinates,
+            last_updated,
+            confirmed=timelines.get("confirmed").latest or 0,
+            deaths=timelines.get("deaths").latest or 0,
+            recovered=timelines.get("recovered").latest or 0,
+            timelines=timelines
+        )
 
         self.state = state
         self.county = county

app/services/location/jhu.py

@@ -10,7 +10,7 @@
 
 from ...caches import check_cache, load_cache
 from ...coordinates import Coordinates
-from ...location import TimelinedLocation
+from ...location.jhu import JHULocation
 from ...models import Timeline
 from ...utils import countries
 from ...utils import date as date_util
@@ -172,7 +172,7 @@ async def get_locations():
 
         # Create location (supporting timelines) and append.
         locations.append(
-            TimelinedLocation(
+            JHULocation(
                 # General info.
                 index,
                 location["country"],

tests/test_location.py

@@ -47,7 +47,7 @@ def test_location_class(
     now = datetime.utcnow().isoformat() + "Z"
 
     # Location.
-    location_obj = location.TimelinedLocation(
+    location_obj = location.Location(
         test_id,
         country,
         province,

tests/test_nyt.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from app.location import TimelinedLocation
+from app.location import Location
 from app.location.nyt import NYTLocation
 from app.services.location import nyt
 from tests.conftest import mocked_strptime_isoformat
@@ -23,7 +23,7 @@ async def test_get_locations(mock_client_session):
     serialized_locations = []
     for location in locations:
         assert isinstance(location, NYTLocation)
-        assert isinstance(location, TimelinedLocation)
+        assert isinstance(location, Location)
 
         # Making sure country population is a non-zero value
         assert location.country_population != 0