Simplifying functions and removing subtle bug (heat stage instead of …

…heat type)

tests/fixtures/thermostats.py

@@ -3,6 +3,7 @@
 from thermostat.util.testing import get_data_path
 from thermostat.regression import runtime_regression
 from thermostat.core import Thermostat, CoreDaySet
+from tempfile import TemporaryDirectory
 
 import pandas as pd
 import numpy as np
@@ -84,6 +85,12 @@ def thermostat_type_1(request):
     thermostats = from_csv(get_data_path(request.param))
     return next(thermostats)
 
+@pytest.fixture(scope="session", params=["../data/metadata_type_1_single.csv"])
+def thermostat_type_1_cache(request):
+    with TemporaryDirectory() as tempdir:
+        thermostats = from_csv(get_data_path(request.param), save_cache=True, cache_path=tempdir)
+        return next(thermostats)
+
 @pytest.fixture(scope="session", params=["../data/metadata_type_2_single.csv"])
 def thermostat_type_2(request):
     thermostats = from_csv(get_data_path(request.param))

tests/test_importers.py

@@ -14,6 +14,7 @@
         thermostat_type_1_too_many_minutes,
         thermostat_type_1_zip_bad,
         thermostat_type_1_data_out_of_order,
+        thermostat_type_1_cache,
         )
 
 def test_import_csv(thermostat_type_1):
@@ -34,6 +35,9 @@ def assert_is_series_with_shape(series, shape):
     assert_is_series_with_shape(thermostat_type_1.temperature_in, (35064,))
     assert_is_series_with_shape(thermostat_type_1.temperature_out, (35064,))
 
+def test_import_csv_cache(thermostat_type_1_cache):
+    assert thermostat_type_1_cache is not None
+
 def test_utc_offset(thermostat_type_1_utc, thermostat_type_1_utc_bad):
     assert(normalize_utc_offset("+0") == datetime.timedelta(0))
     assert(normalize_utc_offset("-0") == datetime.timedelta(0))

thermostat/importers.py

@@ -78,7 +78,7 @@ def save_json_cache(index, thermostat_id, station, cache_path=None):
         os.mkdir(directory)
     except OSError as e:
         if e.errno != errno.EEXIST:
-            raise
+            raise e
 
     json_cache = {}
 
@@ -310,35 +310,42 @@ def get_single_thermostat(thermostat_id, zipcode,
                    "which are out of order, missing, or duplicated.".format(thermostat_id))
         raise RuntimeError(message)
 
+    # Export the data from the cache
+    if save_cache:
+        save_json_cache(hourly_index, thermostat_id, station, cache_path)
+
     # load hourly time series values
     temp_in = _create_series(df.temp_in, hourly_index)
 
-    if has_auxiliary(heat_type) and has_emergency(heat_type):
-        auxiliary_heat_runtime = _create_series(df.auxiliary_heat_runtime, hourly_index)
-        emergency_heat_runtime = _create_series(df.emergency_heat_runtime, hourly_index)
-        if auxiliary_heat_runtime.gt(60).any():
-            warnings.warn("For thermostat {}, auxiliary runtime data was larger than 60 minutes"
-                          " for one or more hours, which is impossible. Please check the data file."
-                          .format(thermostat_id))
-            return
-        if emergency_heat_runtime.gt(60).any():
-            warnings.warn("For thermostat {}, emergency runtime data was larger than 60 minutes"
-                          " for one or more hours, which is impossible. Please check the data file."
-                          .format(thermostat_id))
-            return
-    else:
-        auxiliary_heat_runtime = None
-        emergency_heat_runtime = None
-
     utc_offset = normalize_utc_offset(utc_offset)
     temp_out = get_indexed_temperatures_eeweather(station, hourly_index_utc - utc_offset)
     temp_out.index = hourly_index
 
-    # Export the data from the cache
-    if save_cache:
-        save_json_cache(hourly_index, thermostat_id, station, cache_path)
-
     # load daily time series values
+    auxiliary_heat_runtime, emergency_heat_runtime = _calculate_aux_emerg_runtime(df, thermostat_id, heat_type, heat_stage, hourly_index)
+    cool_runtime = _calculate_cool_runtime(df, thermostat_id, cool_type, cool_stage, hourly_index)
+    heat_runtime = _calculate_heat_runtime(df, thermostat_id, heat_type, heat_stage, hourly_index)
+
+    # create thermostat instance
+    thermostat = Thermostat(
+        thermostat_id,
+        heat_type,
+        heat_stage,
+        cool_type,
+        cool_stage,
+        zipcode,
+        station,
+        temp_in,
+        temp_out,
+        cool_runtime,
+        heat_runtime,
+        auxiliary_heat_runtime,
+        emergency_heat_runtime
+    )
+    return thermostat
+
+
+def _calculate_cool_runtime(df, thermostat_id, cool_type, cool_stage, hourly_index):
     if has_cooling(cool_type):
         cool_runtime_stg1 = df.cool_runtime_stg1
         cool_runtime_stg2 = df.cool_runtime_stg2
@@ -361,13 +368,16 @@ def get_single_thermostat(thermostat_id, zipcode,
                         'Skipping import of thermostat with staged cooling runtime '
                         'greater than cooling equivalent runtime. (id={})'.format(thermostat_id))
                 return
-            df['cool_runtime'] = cool_runtime_both_stg
+
+            cool_runtime = _create_series(cool_runtime_both_stg, hourly_index)
         else:
-            df['cool_runtime'] = cool_runtime_stg1
-        cool_runtime = _create_series(df.cool_runtime, hourly_index)
+            cool_runtime = _create_series(cool_runtime_stg1, hourly_index)
     else:
         cool_runtime = None
+    return cool_runtime
 
+
+def _calculate_heat_runtime(df, thermostat_id, heat_type, heat_stage, hourly_index):
     if has_heating(heat_type):
         heat_runtime_stg1 = df.heat_runtime_stg1
         heat_runtime_stg2 = df.heat_runtime_stg2
@@ -381,38 +391,40 @@ def get_single_thermostat(thermostat_id, zipcode,
                           .format(thermostat_id))
             return
 
-        if has_two_stage_heating(heat_type):
+        if has_two_stage_heating(heat_stage):
             heat_runtime_both_stg = (first_stage_capacity_fraction(heat_type) * heat_runtime_stg1) + heat_runtime_stg2
             heat_runtime_in_bounds = heat_runtime_both_stg.dropna() <= df.heat_equiv_runtime.dropna()
             if not(heat_runtime_in_bounds.all()):
                 warnings.warn(
                         'Skipping import of thermostat with staged heating runtime '
                         'greater than heating equivalent runtime. (id={})'.format(thermostat_id))
                 return
-            df['heat_runtime'] = heat_runtime_both_stg
+            heat_runtime = _create_series(heat_runtime_both_stg, hourly_index)
         else:
-            df['heat_runtime'] = heat_runtime_stg1
-        heat_runtime = _create_series(df.heat_runtime, hourly_index)
+            heat_runtime = _create_series(df.heat_runtime_stg1, hourly_index)
     else:
         heat_runtime = None
+    return heat_runtime
 
-    # create thermostat instance
-    thermostat = Thermostat(
-        thermostat_id,
-        heat_type,
-        heat_stage,
-        cool_type,
-        cool_stage,
-        zipcode,
-        station,
-        temp_in,
-        temp_out,
-        cool_runtime,
-        heat_runtime,
-        auxiliary_heat_runtime,
-        emergency_heat_runtime
-    )
-    return thermostat
+
+def _calculate_aux_emerg_runtime(df, thermostat_id, heat_type, heat_stage, hourly_index):
+    if has_auxiliary(heat_type) and has_emergency(heat_type):
+        auxiliary_heat_runtime = _create_series(df.auxiliary_heat_runtime, hourly_index)
+        emergency_heat_runtime = _create_series(df.emergency_heat_runtime, hourly_index)
+        if auxiliary_heat_runtime.gt(60).any():
+            warnings.warn("For thermostat {}, auxiliary runtime data was larger than 60 minutes"
+                          " for one or more hours, which is impossible. Please check the data file."
+                          .format(thermostat_id))
+            return
+        if emergency_heat_runtime.gt(60).any():
+            warnings.warn("For thermostat {}, emergency runtime data was larger than 60 minutes"
+                          " for one or more hours, which is impossible. Please check the data file."
+                          .format(thermostat_id))
+            return
+    else:
+        auxiliary_heat_runtime = None
+        emergency_heat_runtime = None
+    return auxiliary_heat_runtime, emergency_heat_runtime
 
 
 def _create_series(df, index):