Skip to content

Commit

Permalink
2679: set_range() refactor; unit tests (#2844)
Browse files Browse the repository at this point in the history
* Refactoring set_range()

* Adding more tests

* Adding more tests

* Lint

* No default test

* Better var names

---------

Co-authored-by: Phil Dominguez <“[email protected]”>
  • Loading branch information
phildominguez-gsa and Phil Dominguez authored Nov 22, 2023
1 parent 9dc4033 commit da1d340
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 43 deletions.
131 changes: 123 additions & 8 deletions backend/census_historical_migration/test_excel_creation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):
Expand Down

0 comments on commit da1d340

Please sign in to comment.