Refactor type handling in invoices

Previously, there was no consistent method to create new invoice
columns with an explicit type. This has been the source of some
confusion when debugging and writing test cases. As such, several
changes have been added to reduce confusion over the typing of
invoice columns:
- A function (`_create_column`) has been added to `invoice.py`
to provide a consistent way of creating new columns with an explicit type.
- Many fields will now use Pandas' built-in extension types [1],
such as `BooleanDtype`, which allows them to be nullable. This
also enforces typing, as assigning incompatible values to these
columns will raise a type error
- Decimal precision has been increased to 21. This is to allow
`int64` values to be converted to `pyarrow.decimal128(21, 2)`.
This is useful for writing test cases (no longer have to write
`[1.0, 2.0]`), and in case our invoices have balance fields that
resemble integers
- A new test case base class is made with a function to create a
dataframe with standardized column types. This allows creating
invoice dataframes that more closely resemble what our pipeline
will see during execution
- Removed redundant code in `_get_test_invoice()` of many test classes
- `PIInvoice` will cast the dataframe to `StringDtype` to prevent
TypeError with the new Pandas extension types
[1] https://pandas.pydata.org/docs/user_guide/basics.html#dtypes

Author: QuanMPhm

process_report/invoices/invoice.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 import pandas
+import pyarrow
 
 import process_report.util as util
 
@@ -58,6 +59,12 @@
 CLUSTER_NAME_FIELD = "Cluster Name"
 ###
 
+### Invoice column types
+BALANCE_FIELD_TYPE = pandas.ArrowDtype(pyarrow.decimal128(21, 2))
+STRING_FIELD_TYPE = pandas.StringDtype()
+BOOL_FIELD_TYPE = pandas.BooleanDtype()
+###
+
 
 @dataclass
 class Invoice:
@@ -110,6 +117,11 @@ def _prepare_export(self):
         that should or should not be exported after processing."""
         pass
 
+    def _create_column(self, column_name, column_type, value):
+        """Creates a new column in `self.data` with type and value."""
+        self.data[column_name] = value
+        self.data[column_name] = self.data[column_name].astype(column_type)
+
     def _filter_columns(self):
         """Filters and renames columns before exporting"""
         self.export_data = self.export_data[self.export_columns_list].rename(

process_report/invoices/pi_specific_invoice.py

@@ -105,6 +105,7 @@ def _get_pi_dataframe(self, data, pi):
                     lambda data: data if pandas.isna(data) else f"${data}"
                 )
 
+        pi_projects = pi_projects.astype(invoice.STRING_FIELD_TYPE)
         pi_projects.fillna("", inplace=True)
 
         return pi_projects

process_report/process_report.py

@@ -6,11 +6,11 @@
 import os
 
 import pandas
-import pyarrow
 from nerc_rates import load_from_url
 
 from process_report import util
 from process_report.invoices import (
+    invoice,
     lenovo_invoice,
     nonbillable_invoice,
     billable_invoice,
@@ -420,8 +420,8 @@ def merge_csv(files):
         dataframe = pandas.read_csv(
             file,
             dtype={
-                COST_FIELD: pandas.ArrowDtype(pyarrow.decimal128(12, 2)),
-                RATE_FIELD: str,
+                COST_FIELD: invoice.BALANCE_FIELD_TYPE,
+                RATE_FIELD: invoice.STRING_FIELD_TYPE,
             },
         )
         dataframes.append(dataframe)

process_report/processors/bu_subsidy_processor.py

@@ -19,7 +19,16 @@ def get_project(row):
             else:
                 return project_alloc[: project_alloc.rfind("-")]
 
-        self.data[invoice.PROJECT_NAME_FIELD] = self.data.apply(get_project, axis=1)
+        self._create_column(
+            invoice.PROJECT_NAME_FIELD,
+            invoice.STRING_FIELD_TYPE,
+            self.data.apply(get_project, axis=1),
+        )
+        self._create_column(
+            invoice.SUBSIDY_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            0,
+        )
         self.data[invoice.SUBSIDY_FIELD] = Decimal(0)
 
     def _process(self):

process_report/processors/lenovo_processor.py

@@ -10,7 +10,13 @@ class LenovoProcessor(processor.Processor):
     SU_CHARGE_MULTIPLIER = 1
 
     def _process(self):
-        self.data[invoice.SU_CHARGE_FIELD] = self.SU_CHARGE_MULTIPLIER
-        self.data[invoice.LENOVO_CHARGE_FIELD] = (
-            self.data[invoice.SU_HOURS_FIELD] * self.data[invoice.SU_CHARGE_FIELD]
+        self._create_column(
+            invoice.SU_CHARGE_FIELD,
+            int,
+            self.SU_CHARGE_MULTIPLIER,
+        )
+        self._create_column(
+            invoice.LENOVO_CHARGE_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            self.data[invoice.SU_HOURS_FIELD] * self.data[invoice.SU_CHARGE_FIELD],
         )

process_report/processors/new_pi_credit_processor.py

@@ -183,10 +183,26 @@ def get_initial_credit_amount(
         return (data, old_pi_df)
 
     def _prepare(self):
-        self.data[invoice.CREDIT_FIELD] = None
-        self.data[invoice.CREDIT_CODE_FIELD] = None
-        self.data[invoice.PI_BALANCE_FIELD] = self.data[invoice.COST_FIELD]
-        self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD]
+        self._create_column(
+            invoice.CREDIT_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.CREDIT_CODE_FIELD,
+            invoice.STRING_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.PI_BALANCE_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            self.data[invoice.COST_FIELD],
+        )
+        self._create_column(
+            invoice.BALANCE_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            self.data[invoice.COST_FIELD],
+        )
         self.old_pi_df = self._load_old_pis(self.old_pi_filepath)
 
     def _process(self):

process_report/processors/prepayment_processor.py

@@ -38,11 +38,31 @@ def _load_prepay_debits(prepay_debits_filepath):
         return prepay_debits
 
     def _prepare(self):
-        self.data[invoice.GROUP_NAME_FIELD] = None
-        self.data[invoice.GROUP_INSTITUTION_FIELD] = None
-        self.data[invoice.GROUP_MANAGED_FIELD] = None
-        self.data[invoice.GROUP_BALANCE_FIELD] = None
-        self.data[invoice.GROUP_BALANCE_USED_FIELD] = None
+        self._create_column(
+            invoice.GROUP_NAME_FIELD,
+            invoice.STRING_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.GROUP_INSTITUTION_FIELD,
+            invoice.STRING_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.GROUP_MANAGED_FIELD,
+            invoice.BOOL_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.GROUP_BALANCE_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            None,
+        )
+        self._create_column(
+            invoice.GROUP_BALANCE_USED_FIELD,
+            invoice.BALANCE_FIELD_TYPE,
+            None,
+        )
 
         self.prepay_debits = self._load_prepay_debits(self.prepay_debits_filepath)
         self.group_info_dict = self._get_prepay_group_dict()

process_report/processors/validate_billable_pi_processor.py

@@ -56,7 +56,15 @@ def _str_to_lowercase(data):
         )
 
     def _process(self):
-        self.data[invoice.IS_BILLABLE_FIELD] = self._get_billables(
-            self.data, self.nonbillable_pis, self.nonbillable_projects
+        self._create_column(
+            invoice.IS_BILLABLE_FIELD,
+            invoice.BOOL_FIELD_TYPE,
+            self._get_billables(
+                self.data, self.nonbillable_pis, self.nonbillable_projects
+            ),
+        )
+        self._create_column(
+            invoice.MISSING_PI_FIELD,
+            invoice.BOOL_FIELD_TYPE,
+            self._validate_pi_names(self.data),
         )
-        self.data[invoice.MISSING_PI_FIELD] = self._validate_pi_names(self.data)

process_report/tests/base.py

@@ -3,8 +3,48 @@
 from pathlib import Path
 from unittest import TestCase
 
+import pandas
+import pyarrow
 
-class BaseTestCaseWithTempDir(TestCase):
+
+### Invoice column types
+BOOL_FIELD_TYPE = pandas.BooleanDtype()
+STRING_FIELD_TYPE = pandas.StringDtype()
+BALANCE_FIELD_TYPE = pandas.ArrowDtype(pyarrow.decimal128(21, 2))
+###
+
+
+class BaseTestCase(TestCase):
+    field_type_mapping = {
+        "Is Billable": BOOL_FIELD_TYPE,
+        "Missing PI": BOOL_FIELD_TYPE,
+        "Rate": STRING_FIELD_TYPE,
+        "Cost": BALANCE_FIELD_TYPE,
+        "Credit": BALANCE_FIELD_TYPE,
+        "Subsidy": BALANCE_FIELD_TYPE,
+        "Balance": BALANCE_FIELD_TYPE,
+        "Prepaid Group Managed": BOOL_FIELD_TYPE,
+        "Prepaid Group Balance": BALANCE_FIELD_TYPE,
+        "Prepaid Group Used": BALANCE_FIELD_TYPE,
+    }
+
+    def _create_test_invoice(self, data: dict) -> pandas.DataFrame:
+        """
+        Given a dictionary of data, create a DataFrame that represents an invoice.
+        Also standardizes the type for some columns mostly those that expect monetary
+        values, like BALANCE_FIELD, COST_FIELD, etc. These columns are usually
+        converted to pyarrow.decimal128(21, 2)
+        """
+        standard_dtypes = {}
+        for column_name in data.keys():
+            if column_name in self.field_type_mapping:
+                standard_dtypes[column_name] = self.field_type_mapping[column_name]
+
+        test_invoice = pandas.DataFrame(data)
+        return test_invoice.astype(standard_dtypes)
+
+
+class BaseTestCaseWithTempDir(BaseTestCase):
     def setUp(self):
         self.tempdir = Path(tempfile.TemporaryDirectory(delete=False).name)
 

process_report/tests/unit/invoices/test_base_invoice.py

@@ -1,10 +1,10 @@
 from unittest import TestCase, mock
 import pandas
 
-from process_report.tests import util as test_utils
+from process_report.tests import base, util as test_utils
 
 
-class TestBaseInvoice(TestCase):
+class TestBaseInvoice(base.BaseTestCase):
     def test_filter_exported_columns(self):
         test_invoice = pandas.DataFrame(columns=["C1", "C2", "C3", "C4", "C5"])
         answer_invoice = pandas.DataFrame(columns=["C1", "C3R", "C5R"])
@@ -17,6 +17,20 @@ def test_filter_exported_columns(self):
 
         self.assertTrue(result_invoice.equals(answer_invoice))
 
+    def test_create_column(self):
+        test_invoice = self._create_test_invoice({"Col1": [1, 2], "Col2": [3, 4]})
+        answer_invoice = test_invoice.copy()
+        answer_invoice["Col3"] = "default"
+        answer_invoice["Col4"] = 10.0
+        answer_invoice["Col3"] = answer_invoice["Col3"].astype(base.STRING_FIELD_TYPE)
+        answer_invoice["Col4"] = answer_invoice["Col4"].astype(base.BALANCE_FIELD_TYPE)
+
+        inv = test_utils.new_base_invoice(data=test_invoice)
+        inv._create_column("Col3", base.STRING_FIELD_TYPE, "default")
+        inv._create_column("Col4", base.BALANCE_FIELD_TYPE, 10.0)
+
+        self.assertTrue(inv.data.equals(answer_invoice))
+
 
 class TestUploadToS3(TestCase):
     @mock.patch("process_report.util.get_invoice_bucket")