fix: ValueError, #878

lisphilar · Jul 11, 2021 · e71885e · e71885e
1 parent 9a19bdb
commit e71885e
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/covsirphy/analysis/data_handler.py b/covsirphy/analysis/data_handler.py
@@ -49,6 +49,10 @@ def __init__(self, country, province=None, **kwargs):
         self._area_dict = {"country": str(country), "province": str(province or self.UNKNOWN)}
         # Data {str: instance}
         self._data_dict = dict.fromkeys(self.MAIN_DICT.keys(), None)
+        # Raw data without complement
+        self._subset_df_raw = None
+        # With complement
+        self._subset_df = None
         # Population
         self._population = None
         # Auto complement: manually changed with DataHandler.switch_complement()
@@ -218,15 +222,16 @@ def records_main(self):
         if jhu_data is None:
             raise NotRegisteredMainError(".register(jhu_data)")
         # Subsetting
-        df, self._complemented = jhu_data.records(
-            **self._area_dict,
-            start_date=self._first_date, end_date=self._last_date,
-            population=self._population,
-            **self._complement_dict,
-        )
+        if self._subset_df is None:
+            self._subset_df, self._complemented = jhu_data.records(
+                **self._area_dict,
+                start_date=self._first_date, end_date=self._last_date,
+                population=self._population,
+                **self._complement_dict,
+            )
         # Columns which are included in the main dataset except for 'Date'
-        self._main_cols = list(set(df.columns) - set([self.DATE]))
-        return df
+        self._main_cols = list(set(self._subset_df.columns) - set([self.DATE]))
+        return self._subset_df
 
     def switch_complement(self, whether=None, **kwargs):
         """