Apply decorators to location

app/location/__init__.py

@@ -74,6 +74,32 @@ def serialize(self):
         }
 
 
+def decoratedSerialize(func, inc_timelines, *args, **kwargs):
+    serialized = func()
+    if kwargs.get("timelines") and inc_timelines:
+        # Whether to include the timelines or not.
+        timelines = kwargs["timelines"]
+        if timelines:
+            serialized.update(
+                {
+                    "timelines": {
+                        # Serialize all the timelines.
+                        key: value.serialize()
+                        for (key, value) in self.timelines.items()
+                    }
+                }
+            )
+
+    if kwargs.get("state"):
+        state = kwargs["state"]
+        serialized.update({"state": state})
+
+    if kwargs.get("county"):
+        county = kwargs["county"]
+        serialized.update({"county": county})
+
+    return serialized
+
 class TimelinedLocation(Location):
     """
     A location with timelines.
@@ -96,29 +122,19 @@ def __init__(self, id, country, province, coordinates, last_updated, timelines):
 
         # Set timelines.
         self.timelines = timelines
+        self.inc_timelines = false
 
     # pylint: disable=arguments-differ
-    def serialize(self, timelines=False):
+    @decoratedSerialize(timelines=self.timelines, inc_timelines=self.inc_timelines)
+    def serialize(self, inc_timelines=false):
         """
         Serializes the location into a dict.
-
-        :param timelines: Whether to include the timelines.
+        :param inc_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(
-                {
-                    "timelines": {
-                        # Serialize all the timelines.
-                        key: value.serialize()
-                        for (key, value) in self.timelines.items()
-                    }
-                }
-            )
+        self.inc_timelines = inc_timelines
 
         # Return the serialized location.
         return serialized

app/location/csbs.py

@@ -25,19 +25,16 @@ def __init__(self, id, state, county, coordinates, last_updated, confirmed, deat
         self.state = state
         self.county = county
 
-    def serialize(self, timelines=False):  # pylint: disable=arguments-differ,unused-argument
+    @decoratedSerialize(state=self.state, county=self.county)
+    def serialize(self):  # pylint: disable=arguments-differ,unused-argument
         """
         Serializes the location into a dict.
+
 
         :returns: The serialized location.
         :rtype: dict
         """
         serialized = super().serialize()
 
-        # Update with new fields.
-        serialized.update(
-            {"state": self.state, "county": self.county,}
-        )
-
         # Return the serialized location.
         return serialized

app/location/nyt.py

@@ -14,19 +14,14 @@ def __init__(self, id, state, county, coordinates, last_updated, timelines):
         self.state = state
         self.county = county
 
-    def serialize(self, timelines=False):  # pylint: disable=arguments-differ,unused-argument
+    @decoratedSerialize(timelines=self.timelines, state=self.state, county=self.county)
+    def serialize(self, inc_timelines=False):  # pylint: disable=arguments-differ,unused-argument
         """
         Serializes the location into a dict.
-
         :returns: The serialized location.
         :rtype: dict
         """
-        serialized = super().serialize(timelines)
-
-        # Update with new fields.
-        serialized.update(
-            {"state": self.state, "county": self.county,}
-        )
+        serialized = super().serialize(inc_timelines)
 
         # Return the serialized location.
         return serialized