Update cpi.py

economics/cpi.py

@@ -19,6 +19,8 @@
 import datetime
 import collections
 import requests
+import pickle
+import os
 
 CPIResult = collections.namedtuple('CPI', 'date value')
 
@@ -28,28 +30,69 @@ class CPI:
     Provides a Pythonic interface to Consumer Price Index data packages
     """
 
-    def __init__(self, country="all"):
+    def __init__(self, country="all", data_source=None):
         """
         Initialise a CPI instance.
+
+        Parameters:
+        - country (str): the country to retrieve data for. Defaults to "all".
+        - data_source (str): the data source to use. Currently supported sources are
+        "world_bank" (default) and "local_cache".
         """
 
         # Initialise empty data structures
         self.data = collections.ChainMap()
 
+        # Set data source
+        if data_source:
+            self.data_source = data_source
+        else:
+            self.data_source = "world_bank"
+
         # Load the data into the data structures
         self.load(country)
         self.country = country
 
     def load(self, country="all"):
         """
         Load data
+
+        Parameters:
+        - country (str): the country to retrieve data for. Defaults to "all".
         """
-        u = ("https://api.worldbank.org/v2/countries/{}/indicators/CPTOTSAXN"
-             "?format=json&per_page=10000").format(country)
-        r = requests.get(u)
-        j = r.json()
-        cpi_data = j[1]
+        if self.data_source == "world_bank":
+            u = ("https://api.worldbank.org/v2/countries/{}/indicators/CPTOTSAXN"
+                 "?format=json&per_page=10000").format(country)
+            r = requests.get(u)
+            try:
+                j = r.json()
+            except ValueError:
+                print("Error parsing data from World Bank API")
+                return
+            try:
+                cpi_data = j[1]
+            except IndexError:
+                print("Error: No data returned from World Bank API")
+                return
+            self._store_data(cpi_data)
+        elif self.data_source == "local_cache":
+            try:
+                with open("cpi_data.pkl", "rb") as f:
+                    cpi_data = pickle.load(f)
+            except FileNotFoundError:
+                print("Error: CPI data cache not found")
+                return
+            self._store_data(cpi_data)
+        else:
+            raise ValueError("Unrecognized data source: {}".format(self.data_source))
 
+    def _store_data(self, cpi_data):
+        """
+        Store CPI data in the data attribute
+
+        Parameters:
+        - cpi_data (list): a list of dictionaries containing the CPI data
+        """
         # Loop through the rows of the datapackage with the help of data
         for row in cpi_data:
             # Get the code and the name and transform to uppercase
@@ -58,7 +101,7 @@ def load(self, country="all"):
             iso_2 = row["country"]["id"].upper()
             name = row["country"]['value'].upper()
             # Get the date (which is in the field Year) and the CPI value
-            date = row['date']
+            date = row['date '] # Note the trailing space in the field name
             cpi = row['value']
             for key in [iso_3, iso_2, name]:
                 existing = self.data.get(key, {})
@@ -68,8 +111,19 @@ def load(self, country="all"):
 
     def get(self, date=datetime.date.today(), country=None):
         """
-        Get the CPI value for a specific time. Defaults to today. This uses
-        the closest method internally but sets limit to one day.
+        Get the CPI value for a specific time. Defaults to today.
+
+        Parameters:
+        - date (datetime.date or int or str): the date to retrieve the CPI value for.
+        If not provided, defaults to today.
+        - country (str): the country to retrieve data for. Defaults to the country
+        specified when the CPI instance was created.
+
+        Returns:
+        - CPIResult: a named tuple with the date and value.
+
+        Raises:
+        - ValueError: if the data is not available for the specified date and country.
         """
         if not country:
             country = self.country
@@ -83,4 +137,24 @@ def get(self, date=datetime.date.today(), country=None):
             raise ValueError("Missing CPI data for {} for {}".format(
                 country, date))
 
+        return CPIResult(date=date, value=cpi)
+
+    def save_cache(self):
+        """
+        Save the current data to a local cache file
+        """
+        with open("cpi_data.pkl", "wb") as f:
+            pickle.dump(self.data.values(), f)
+
+        def clear_cache(self):
+        """
+        Delete the local cache file
+        """
+        try:
+            os.remove("cpi_data.pkl")
+        except FileNotFoundError:
+            pass
+
+
+
         return CPIResult(date=date, value=cpi)