Fix export issues

apps/ary/export/__init__.py

@@ -104,6 +104,7 @@ def normalize_assessment(assessment_export_data):
     v.update(common_defaults)
 
 
+# NOTE: This is magic function, but make it simpler
 def add_assessment_to_rows(sheets, assessment):
     """
     sheets = {
@@ -146,8 +147,11 @@ def add_new_keys(keys, data, default=None):
             # If columns data is empty, add new data to account for empty row
             # assessment data is then appended
             if not columns_data:
-                ass_sample = assessment_col_data[0]\
-                    if isinstance(assessment_col_data, list) else assessment_col_data
+                if isinstance(assessment_col_data, list):
+                    # TODO: Try to check if it should be dict or None
+                    ass_sample = assessment_col_data[0] if assessment_col_data else {}
+                else:
+                    ass_sample = assessment_col_data
                 if isinstance(ass_sample, dict):
                     columns_data = [{}]
                 else:
@@ -189,6 +193,6 @@ def add_new_keys(keys, data, default=None):
 
 def get_export_data_for_assessments(assessments):
     if not assessments:
-        return None
+        return {}
     data = normalize_assessment(get_export_data(assessments[0]))
     return reduce(add_assessment_to_rows, assessments[1:], data)

apps/ary/export/locations_info.py

@@ -19,7 +19,9 @@ def get_locations_info(assessment):
     region_geos = {x['key']: x for x in region.geo_options}
 
     for area in geo_areas:
-        geo_info = region_geos[str(area.id)]
+        geo_info = region_geos.get(str(area.id))
+        if geo_info is None:
+            continue
         level = geo_info['admin_level']
         key = f'Admin {level}'
 

apps/export/assessments/excel_exporter.py

@@ -213,7 +213,7 @@ def __init__(self, sheets_data):
         }
         """
         self.wb = WorkBook()
-        self.sheets_data = sheets_data
+        self.sheets_data = sheets_data or {}
         self.wb_sheets = {}
         self.sheet_headers = {}
 
@@ -285,8 +285,9 @@ def export(self):
         # Add data rows
         self.add_data_rows_to_sheets()
 
-        # Remove default sheet
-        self.wb.wb.remove(self.wb.wb.get_sheet_by_name('Sheet'))
+        # Remove default sheet only if other sheets present
+        if self.wb_sheets:
+            self.wb.wb.remove(self.wb.wb.get_sheet_by_name('Sheet'))
 
         buffer = self.wb.save()
         filename = generate_filename('Assessments Export', 'xlsx')