diff --git a/backend/census_historical_migration/test_excel_creation.py b/backend/census_historical_migration/test_excel_creation.py index fca6a1e77d..8d7f9c8bf3 100644 --- a/backend/census_historical_migration/test_excel_creation.py +++ b/backend/census_historical_migration/test_excel_creation.py @@ -9,22 +9,137 @@ class ExcelCreationTests(TestCase): - range_name, start_val, end_val, cell = "my_range", "foo", "bar", "A6" + range_name = "my_range" - def test_set_range(self): + def init_named_range(self, coord): """ - Standard use case + Create and return a workbook with a named range for the given coordinate """ wb = Workbook() ws = wb.active # Create named range - ref = f"{quote_sheetname(ws.title)}!{absolute_coordinate(self.cell)}" + ref = f"{quote_sheetname(ws.title)}!{absolute_coordinate(coord)}" + defn = DefinedName(self.range_name, attr_text=ref) + wb.defined_names.add(defn) + + return wb + + def test_set_range_standard(self): + """ + Standard use case + """ + wb = self.init_named_range("A6") + ws = wb.active + + ws.cell(row=6, column=1, value="foo") + self.assertEqual(ws["A6"].value, "foo") + + set_range(wb, self.range_name, ["bar"]) + self.assertEqual(ws["A6"].value, "bar") + + def test_set_range_default(self): + """ + Using a default value + """ + wb = self.init_named_range("A6") + ws = wb.active + + ws.cell(row=6, column=1, value="foo") + self.assertEqual(ws["A6"].value, "foo") + + set_range(wb, self.range_name, [None], default="bar") + self.assertEqual(ws["A6"].value, "bar") + + def test_set_range_no_default(self): + """ + Default to empty string when no value or default given + """ + wb = self.init_named_range("A6") + ws = wb.active + + ws.cell(row=6, column=1, value="foo") + self.assertEqual(ws["A6"].value, "foo") + + set_range( + wb, + self.range_name, + [None], + ) + self.assertEqual(ws["A6"].value, "") + + def test_set_range_conversion(self): + """ + Applies given conversion function + """ + wb = self.init_named_range("A6") + ws = wb.active + + ws.cell(row=6, column=1, value="foo") + self.assertEqual(ws["A6"].value, "foo") + + set_range(wb, self.range_name, ["1"], None, int) + self.assertEqual(ws["A6"].value, 1) # str -> int + + def test_set_range_multiple_values(self): + """ + Setting multiple values + """ + wb = self.init_named_range("A1:A2") + ws = wb.active + + ws.cell(row=1, column=1, value=0) + self.assertEqual(ws["A1"].value, 0) + + ws.cell(row=2, column=1, value=0) + self.assertEqual(ws["A2"].value, 0) + + set_range(wb, self.range_name, ["1", "2"]) + self.assertEqual(ws["A1"].value, "1") + self.assertEqual(ws["A2"].value, "2") + + def test_set_range_fewer_values(self): + """ + Number of values is less than the range size + """ + wb = self.init_named_range("A1:A2") + ws = wb.active + + ws.cell(row=1, column=1, value="foo") + self.assertEqual(ws["A1"].value, "foo") + + ws.cell(row=2, column=1, value="foo") + self.assertEqual(ws["A2"].value, "foo") + + set_range(wb, self.range_name, ["bar"]) + self.assertEqual(ws["A1"].value, "bar") # New value + self.assertEqual(ws["A2"].value, "foo") # Unchanged + + def test_set_range_multi_dests(self): + """ + Error when multiple destinations found + """ + wb = Workbook() + ws = wb.active + + # Create named range with multiple destinations + ref = f"{quote_sheetname(ws.title)}!{absolute_coordinate('A6')}, \ + {quote_sheetname(ws.title)}!{absolute_coordinate('B6')}" defn = DefinedName(self.range_name, attr_text=ref) wb.defined_names.add(defn) - ws.cell(row=6, column=1, value=self.start_val) - self.assertEqual(ws[self.cell].value, self.start_val) + self.assertEqual(len(list(wb.defined_names[self.range_name].destinations)), 2) + self.assertRaises(ValueError, set_range, wb, self.range_name, ["bar"]) + + def test_set_range_ws_missing(self): + """ + Error when the named range isn't in the given WS + """ + wb = Workbook() + + # Create named range with bad sheet title + ref = f"{quote_sheetname('wrong name')}!{absolute_coordinate('A6')}" + defn = DefinedName(self.range_name, attr_text=ref) + wb.defined_names.add(defn) - set_range(wb, self.range_name, [self.end_val]) - self.assertEqual(ws[self.cell].value, self.end_val) + self.assertRaises(KeyError, set_range, wb, self.range_name, ["bar"]) diff --git a/backend/census_historical_migration/workbooklib/excel_creation_utils.py b/backend/census_historical_migration/workbooklib/excel_creation_utils.py index 58a2443442..c5a14c9023 100644 --- a/backend/census_historical_migration/workbooklib/excel_creation_utils.py +++ b/backend/census_historical_migration/workbooklib/excel_creation_utils.py @@ -1,12 +1,16 @@ from census_historical_migration.exception_utils import DataMigrationError from census_historical_migration.base_field_maps import WorkbookFieldInDissem +from census_historical_migration.workbooklib.templates import sections_to_template_paths from census_historical_migration.sac_general_lib.report_id_generator import ( dbkey_to_report_id, ) -from census_historical_migration.workbooklib.templates import sections_to_template_paths -from openpyxl.utils.cell import column_index_from_string from playhouse.shortcuts import model_to_dict +from openpyxl.utils.cell import ( + rows_from_range, + coordinate_from_string, + column_index_from_string, +) import sys import logging @@ -15,42 +19,51 @@ logger = logging.getLogger(__name__) -# Helper to set a range of values. -# Takes a named range, and then walks down the range, -# filling in values from the list past in (values). def set_range(wb, range_name, values, default=None, conversion_fun=str): + """ + Helper to set a range of values. Takes a named range, and then walks down + the range, filling in the given values. + + wb (Workbook) The workbook + range_name (string) Name of the range to set + values (iterable) Values to set within the range + default (any) Default value to use; defaults to None. + conversion (func) Conversion function to apply to individual values; defaults to str(). + """ the_range = wb.defined_names[range_name] - dest = list(the_range.destinations)[0] - sheet_title = dest[0] - ws = wb[sheet_title] - - start_cell = dest[1].replace("$", "").split(":")[0] - col = column_index_from_string(start_cell[0]) - start_row = int(start_cell[1]) - - for ndx, v in enumerate(values): - row = ndx + start_row - if v: - # This is a very noisy statement, showing everything - # written into the workbook. - # print(f'{range_name} c[{row}][{col}] <- {type(v)} len({len(v)}) {default}') - if v is not None: - ws.cell(row=row, column=col, value=conversion_fun(v)) - if len(str(v)) == 0 and default is not None: - # This is less noisy. Shows up for things like - # empty findings counts. 2023 submissions - # require that field to be 0, not empty, - # if there are no findings. - # print('Applying default') - ws.cell(row=row, column=col, value=conversion_fun(default)) - if not v: - if default is not None: - ws.cell(row=row, column=col, value=conversion_fun(default)) - else: - ws.cell(row=row, column=col, value="") + dests = the_range.destinations + + sheet_title, coord = None, None + for cur_sheet_title, cur_coord in dests: + if sheet_title or coord: + # `destinations` is meant to be iterated over, but we only expect one value + raise ValueError(f"{range_name} has more than one destination") else: - # Leave it blank if we have no default passed in - pass + sheet_title, coord = cur_sheet_title, cur_coord + + ws = None + try: + ws = wb[sheet_title] + except KeyError: + raise KeyError(f"Sheet title '{sheet_title}' not found in workbook") + + values = list(values) + for i, row in enumerate(rows_from_range(coord)): + # Iterate over the rows, but stop when we run out of values + value = None + try: + value = values[i] + except IndexError: + break + + # Get the row and column to set the current value + cell = row[0] # [('B12',)] -> ('B12',) + col_str, row = coordinate_from_string(cell) # ('B12',) -> 'B', 12 + col = column_index_from_string(col_str) # 'B' -> 2 + + # Set the value of the cell + converted_value = conversion_fun(value) if value else default or "" + ws.cell(row=row, column=col, value=converted_value) def set_uei(Gen, wb, dbkey):