Make Population a Property of Location Object

.gitignore

@@ -63,4 +63,7 @@ coverage.xml
 docs/_build/
 
 # PyBuilder
-target/
+target/
+
+# OSX Stuff
+.DS_Store

README.md

@@ -128,6 +128,7 @@ __Sample response__
     "id": 39,
     "country": "Norway",
     "country_code": "NO",
+    "country_population": 5009150,
     "province": "",
     "county": "",
     "last_updated": "2020-03-21T06:59:11.315422Z",
@@ -173,6 +174,7 @@ __Sample response__
       "id": 0,
       "country": "Thailand",
       "country_code": "TH",
+      "country_population": 67089500,
       "province": "",
       "county": "",
       "last_updated": "2020-03-21T06:59:11.315422Z",
@@ -251,6 +253,7 @@ __Sample Response__
       "id": 16,
       "country": "Italy",
       "country_code": "IT",
+      "country_population": 60340328,
       "province": "",
       "county": "",
       "last_updated": "2020-03-23T13:32:23.913872Z",
@@ -290,6 +293,7 @@ __Sample Response__
       "id": 0,
       "country": "US",
       "country_code": "US",
+      "country_population": 310232863,
       "province": "New York",
       "state": "New York",
       "county": "New York",
@@ -308,6 +312,7 @@ __Sample Response__
       "id": 1,
       "country": "US",
       "country_code": "US",
+      "country_population": 310232863,
       "province": "New York",
       "state": "New York",
       "county": "Westchester",

app/location/__init__.py

@@ -1,5 +1,6 @@
 from ..coordinates import Coordinates
 from ..utils import countrycodes
+from ..utils.populations import country_population
 
 class Location:
     """
@@ -25,9 +26,22 @@ def __init__(self, id, country, province, coordinates, last_updated, confirmed,
     def country_code(self):
         """
         Gets the alpha-2 code represention of the country. Returns 'XX' if none is found.
+
+        :returns: The country code.
+        :rtype: str
         """
         return (countrycodes.country_code(self.country) or countrycodes.default_code).upper()
 
+    @property
+    def country_population(self):
+        """
+        Gets the population of this location.
+
+        :returns: The population.
+        :rtype: int
+        """
+        return country_population(self.country_code)
+
     def serialize(self):
         """
         Serializes the location into a dict.
@@ -37,10 +51,11 @@ def serialize(self):
         """
         return {
             # General info.
-            'id'          : self.id,
-            'country'     : self.country, 
-            'country_code': self.country_code,
-            'province'    : self.province,
+            'id'                : self.id,
+            'country'           : self.country, 
+            'country_code'      : self.country_code,
+            'country_population': self.country_population,
+            'province'          : self.province,
 
             # Coordinates.
             'coordinates': self.coordinates.serialize(),
@@ -93,4 +108,4 @@ def serialize(self, timelines = False):
             }})
 
         # Return the serialized location.
-        return serialized
+        return serialized

app/models/location.py

@@ -10,6 +10,7 @@ class Location(BaseModel):
     id: int
     country: str
     country_code: str
+    country_population: int = None
     county: str = ''
     province: str = ''
     last_updated: str # TODO use datetime.datetime type.

app/utils/countrycodes.py

@@ -1,3 +1,5 @@
+from itertools import chain
+
 # Default country code.
 default_code = "XX"
 
@@ -364,12 +366,14 @@ def country_code(country):
     Return two letter country code (Alpha-2) according to https://en.wikipedia.org/wiki/ISO_3166-1
     Defaults to "XX".
     """
+    # Look in synonyms if not found.
+    if not country in is_3166_1 and country in synonyms:
+        country = synonyms[country]
+
+    # Return code if country was found.
     if country in is_3166_1:
         return is_3166_1[country]
-    else:
-        if country in synonyms:
-            synonym = synonyms[country]
-            return is_3166_1[synonym]
-        else:
-            print ("No country_code found for '" + country + "'. Using '" + default_code + "'")
-            return default_code
+
+    # Default to default_code.
+    print ("No country_code found for '" + country + "'. Using '" + default_code + "'")
+    return default_code

app/utils/populations.py

@@ -0,0 +1,43 @@
+import requests
+from io import StringIO, BytesIO
+from cachetools import cached, TTLCache
+from zipfile import ZipFile, ZipInfo
+from .countrycodes import country_code
+
+# Fetching of the populations.
+def fetch_populations():
+    """
+    Returns a dictionary containing the population of each country fetched from the GeoNames (https://www.geonames.org/).
+
+    :returns: The mapping of populations.
+    :rtype: dict
+    """
+    print ("Fetching populations...")
+
+    # Mapping of populations
+    mappings = {}
+
+    # Fetch the countries.
+    countries = requests.get("http://api.geonames.org/countryInfoJSON?username=dperic").json()['geonames']
+
+    # Go through all the countries and perform the mapping.
+    for country in countries:
+        mappings.update({ country["countryCode"]: int(country["population"]) or None })
+
+    # Finally, return the mappings.
+    return mappings
+
+# Mapping of alpha-2 codes country codes to population.
+populations = fetch_populations()
+
+# Retrieving.
+def country_population(country_code, default=None):
+    """
+    Fetches the population of the country with the provided country code.
+
+    :returns: The population.
+    :rtype: int
+    """
+    return populations.get(country_code, default)
+
+

tests/test_location.py

@@ -10,20 +10,20 @@ def __init__(self, latest):
 
     return TestTimeline(args[0])
 
-@pytest.mark.parametrize("test_id, country, country_code, province, latitude, longitude, confirmed_latest, deaths_latest, recovered_latest", [
-    (0, "Thailand", "TH", "", 15, 100, 1000, 1111, 22222),
-    (1, "Deutschland", "DE", "", 15, 100, 1000, 1111, 22222),
-    (2, "Cruise Ship", "XX", "", 15, 100, 1000, 1111, 22222)
+@pytest.mark.parametrize("test_id, country, country_code, country_population, province, latitude, longitude, confirmed_latest, deaths_latest, recovered_latest", [
+    (0, "Thailand", "TH", 1000, "", 15, 100, 1000, 1111, 22222),
+    (1, "Deutschland", "DE", 1000, "", 15, 100, 1000, 1111, 22222),
+    (2, "Cruise Ship", "XX", 1000, "", 15, 100, 1000, 1111, 22222)
 ])
 @mock.patch('app.timeline.Timeline', side_effect=mocked_timeline)
-def test_location_class(mocked_timeline, test_id, country, country_code, province, latitude, longitude, confirmed_latest, deaths_latest, recovered_latest):
+def test_location_class(mocked_timeline, test_id, country, country_code, country_population, province, latitude, longitude, confirmed_latest, deaths_latest, recovered_latest):
 
     # id, country, province, coordinates, confirmed, deaths, recovered
     coords = coordinates.Coordinates(latitude=latitude, longitude=longitude)
 
     # Timelines
     confirmed = timeline.Timeline(confirmed_latest)
-    deaths = timeline.Timeline(deaths_latest)
+    deaths    = timeline.Timeline(deaths_latest)
     recovered = timeline.Timeline(recovered_latest)
 
     # Date now.
@@ -37,23 +37,4 @@ def test_location_class(mocked_timeline, test_id, country, country_code, provinc
     })
 
     assert location_obj.country_code == country_code
-
-    #validate serialize
-    check_dict = {
-        'id': test_id,
-        'country': country,
-        'country_code': country_code,
-        'province': province,
-        'last_updated': now,
-        'coordinates': {
-            'latitude': latitude,
-            'longitude': longitude
-        },
-        'latest': {
-            'confirmed': confirmed_latest,
-            'deaths': deaths_latest,
-            'recovered': recovered_latest
-        }
-    }
-
-    assert location_obj.serialize() == check_dict
+    assert not location_obj.serialize() == None