add 3 rma yield datasets

Author: eschechter

dweather_client/client.py

@@ -12,7 +12,7 @@
 from timezonefinder import TimezoneFinder
 from dweather_client import gridded_datasets
 from dweather_client.storms_datasets import IbtracsDataset, AtcfDataset, SimulatedStormsDataset
-from dweather_client.ipfs_queries import StationDataset, ScoYieldDataset, AemoPowerDataset, AemoGasDataset
+from dweather_client.ipfs_queries import StationDataset, YieldDatasets, AemoPowerDataset, AemoGasDataset
 from dweather_client.ipfs_errors import *
 import ipfshttpclient
 
@@ -216,7 +216,7 @@ def get_station_history(
 
     return history
 
-def get_yield_history(commodity, state, county, ipfs_timeout=None):
+def get_yield_history(commodity, state, county, dataset="sco-yearly", ipfs_timeout=None):
     """
     return:
         string containing yield data in csv format
@@ -228,8 +228,10 @@ def get_yield_history(commodity, state, county, ipfs_timeout=None):
         You can look up code values at:
         https://webapp.rma.usda.gov/apps/RIRS/AreaPlanHistoricalYields.aspx
     """
+    if "imputed" in dataset and commodity != "0081":
+        raise ValueError("Imputed currently only available for soybeans (commodity code 0081)")
     try:
-        return ScoYieldDataset(ipfs_timeout=ipfs_timeout).get_data(commodity, state, county)
+        return YieldDatasets(dataset, ipfs_timeout=ipfs_timeout).get_data(commodity, state, county)
     except ipfshttpclient.exceptions.ErrorResponse:
         raise ValueError("Invalid commodity/state/county code combination")
 

dweather_client/ipfs_queries.py

@@ -385,20 +385,33 @@ def get_data(self, station):
         with gzip.open(self.get_file_object(file_name)) as gz:
             return gz.read().decode('utf-8')
 
-class ScoYieldDataset(IpfsDataset):
+class YieldDatasets(IpfsDataset):
     """
-    Instantiable class used for pulling in "ghcnd" or "ghcnd-imputed-daily" station data
+    Instantiable class used for pulling in sco and rma transitional yield data
     """
     @property
     def dataset(self):
-        return "sco-yearly"
+        return self._dataset
+
+    def __init__(self, dataset, ipfs_timeout=None):
+        if dataset not in {
+            "sco-yearly",
+            "rmasco_imputed-yearly",
+            "rma_t_yield-single-value",
+            "rma_t_yield_imputed-single-value"
+        }:
+            raise ValueError("Invalid yield dataset")
+        super().__init__(ipfs_timeout=ipfs_timeout)
+        self._dataset = dataset
 
     def get_data(self, commodity, state, county):
         super().get_data()
-        file_name = f"{self.head}/{commodity}-{state}-{county}.csv"
+        if "imputed" in self.dataset:
+            file_name = f"{self.head}/{state}-{county}.csv"
+        else:
+            file_name = f"{self.head}/{commodity}-{state}-{county}.csv"
         return self.get_file_object(file_name).read().decode("utf-8")
 
-
 class AemoDataset(IpfsDataset):
     """
     Abstract class from which all AEMO datasets inherit
@@ -431,7 +444,6 @@ def get_data(self):
             ret_dict = {**ret_dict, **new_dict}
         return ret_dict
 
-
 class AemoPowerDataset(AemoDataset):
     """
     Instantiable class for AEMO Victoria power data