diff --git a/hub/management/commands/base_importers.py b/hub/management/commands/base_importers.py index 2850c37a6..bc6854dae 100644 --- a/hub/management/commands/base_importers.py +++ b/hub/management/commands/base_importers.py @@ -1,4 +1,3 @@ -from functools import cache from time import sleep from django.core.management.base import BaseCommand @@ -53,10 +52,15 @@ def delete_data(self): data_type=data_type, area__area_type__code=self.area_type ).delete() - @cache def get_area_type(self): return AreaType.objects.get(code=self.area_type) + def get_cons_col(self): + if hasattr(self, "cons_col_map"): + return self.cons_col_map[self.area_type] + + return self.cons_col + def add_data_sets(self, df=None): for name, config in self.data_sets.items(): label = self.get_label(config) @@ -331,11 +335,11 @@ def set_data_type(self): def get_dataframe(self): df = pd.read_csv(self.data_file) - df = df.astype({self.cons_col: "str"}) + df = df.astype({self.get_cons_col(): "str"}) return df def _get_areas_from_row(self, row): - value = row[self.cons_col] + value = row[self.get_cons_col()] if self.uses_gss: areas = Area.objects.filter(gss__in=value.split(",")) else: diff --git a/hub/management/commands/import_hnh_polling_data.py b/hub/management/commands/import_hnh_polling_data.py index cbf461575..5dc051df1 100644 --- a/hub/management/commands/import_hnh_polling_data.py +++ b/hub/management/commands/import_hnh_polling_data.py @@ -349,6 +349,7 @@ def log(self, message): def extract_and_save_data(self): self.log(self.message) + area_type = self.get_area_type() for file in self.files: self.log(file["defaults"]["label"]) @@ -371,7 +372,7 @@ def extract_and_save_data(self): data_type, created = DataType.objects.update_or_create( data_set=data_set, name=data_type_slug, - area_type=self.get_area_type(), + area_type=area_type, defaults={ "data_type": "percent", "label": col["label"],