Allow invoice subclasses to explicitly declare their exported columns and column names

Two class attributes, `export_columns_list` and `exported_columns_map`,
has been added to `Invoice`, along with a class function `_filter_columns()`.
Subclasses of `Invoice` must now define `export_columns_list`, containing
the ordered list of columns that must be exported in their respective
invoices. Subclasses can optional define `exported_columns_map`,
containing mappings between "internal" column names and what their
name should be when exported.

The field name `RATE_FIELD` has been added to `invoice.py`. It was
previously forgotten.

Author: QuanMPhm

process_report/invoices/NERC_total_invoice.py

@@ -12,6 +12,24 @@ class NERCTotalInvoice(invoice.Invoice):
         "University of Rhode Island",
     ]
 
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PROJECT_FIELD,
+        invoice.PROJECT_ID_FIELD,
+        invoice.PI_FIELD,
+        invoice.INVOICE_EMAIL_FIELD,
+        invoice.INVOICE_ADDRESS_FIELD,
+        invoice.INSTITUTION_FIELD,
+        invoice.INSTITUTION_ID_FIELD,
+        invoice.SU_HOURS_FIELD,
+        invoice.SU_TYPE_FIELD,
+        invoice.RATE_FIELD,
+        invoice.COST_FIELD,
+        invoice.CREDIT_FIELD,
+        invoice.CREDIT_CODE_FIELD,
+        invoice.BALANCE_FIELD,
+    ]
+
     @property
     def output_path(self) -> str:
         return f"NERC-{self.invoice_month}-Total-Invoice.csv"

process_report/invoices/billable_invoice.py

@@ -14,6 +14,24 @@
 class BillableInvoice(invoice.Invoice):
     PI_S3_FILEPATH = "PIs/PI.csv"
 
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PROJECT_FIELD,
+        invoice.PROJECT_ID_FIELD,
+        invoice.PI_FIELD,
+        invoice.INVOICE_EMAIL_FIELD,
+        invoice.INVOICE_ADDRESS_FIELD,
+        invoice.INSTITUTION_FIELD,
+        invoice.INSTITUTION_ID_FIELD,
+        invoice.SU_HOURS_FIELD,
+        invoice.SU_TYPE_FIELD,
+        invoice.RATE_FIELD,
+        invoice.COST_FIELD,
+        invoice.CREDIT_FIELD,
+        invoice.CREDIT_CODE_FIELD,
+        invoice.BALANCE_FIELD,
+    ]
+
     old_pi_filepath: str
     updated_old_pi_df: pandas.DataFrame
 

process_report/invoices/bu_internal_invoice.py

@@ -5,21 +5,23 @@
 
 @dataclass
 class BUInternalInvoice(invoice.Invoice):
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PI_FIELD,
+        "Project",
+        invoice.COST_FIELD,
+        invoice.CREDIT_FIELD,
+        invoice.BU_BALANCE_FIELD,
+        invoice.PI_BALANCE_FIELD,
+    ]
+    exported_columns_map = {
+        invoice.BU_BALANCE_FIELD: "Subsidy",
+        invoice.PI_BALANCE_FIELD: "Balance",
+    }
+
     subsidy_amount: int
 
     def _prepare_export(self):
-        self.data = self.data[
-            [
-                invoice.INVOICE_DATE_FIELD,
-                invoice.PI_FIELD,
-                "Project",
-                invoice.COST_FIELD,
-                invoice.CREDIT_FIELD,
-                invoice.BU_BALANCE_FIELD,
-                invoice.BALANCE_FIELD,
-            ]
-        ]
-
         self.data = self._sum_project_allocations(self.data)
 
     def _sum_project_allocations(self, dataframe):
@@ -28,7 +30,12 @@ def _sum_project_allocations(self, dataframe):
         each unique project, summing up its allocations' costs"""
         project_list = dataframe["Project"].unique()
         data_no_dup = dataframe.drop_duplicates("Project", inplace=False)
-        sum_fields = [invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD]
+        sum_fields = [
+            invoice.COST_FIELD,
+            invoice.CREDIT_FIELD,
+            invoice.BU_BALANCE_FIELD,
+            invoice.PI_BALANCE_FIELD,
+        ]
         for project in project_list:
             project_mask = dataframe["Project"] == project
             no_dup_project_mask = data_no_dup["Project"] == project

process_report/invoices/invoice.py

@@ -23,6 +23,7 @@
 INSTITUTION_ID_FIELD = "Institution - Specific Code"
 SU_HOURS_FIELD = "SU Hours (GBhr or SUhr)"
 SU_TYPE_FIELD = "SU Type"
+RATE_FIELD = "Rate"
 COST_FIELD = "Cost"
 CREDIT_FIELD = "Credit"
 CREDIT_CODE_FIELD = "Credit Code"
@@ -38,6 +39,9 @@
 
 @dataclass
 class Invoice:
+    export_columns_list = list()
+    exported_columns_map = dict()
+
     name: str
     invoice_month: str
     data: pandas.DataFrame
@@ -83,7 +87,14 @@ def _prepare_export(self):
         that should or should not be exported after processing."""
         pass
 
+    def _filter_columns(self):
+        """Filters and renames columns before exporting"""
+        self.data = self.data[self.export_columns_list].rename(
+            columns=self.exported_columns_map
+        )
+
     def export(self):
+        self._filter_columns()
         self.data.to_csv(self.output_path, index=False)
 
     def export_s3(self, s3_bucket):

process_report/invoices/lenovo_invoice.py

@@ -7,17 +7,17 @@
 class LenovoInvoice(invoice.Invoice):
     LENOVO_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
 
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PROJECT_FIELD,
+        invoice.INSTITUTION_FIELD,
+        invoice.SU_HOURS_FIELD,
+        invoice.SU_TYPE_FIELD,
+        "Charge"
+    ]
+    exported_columns_map = {invoice.SU_HOURS_FIELD: "SU Hours"}
+
     def _prepare_export(self):
         self.data = self.data[
             self.data[invoice.SU_TYPE_FIELD].isin(self.LENOVO_SU_TYPES)
-        ][
-            [
-                invoice.INVOICE_DATE_FIELD,
-                invoice.PROJECT_FIELD,
-                invoice.INSTITUTION_FIELD,
-                invoice.SU_HOURS_FIELD,
-                invoice.SU_TYPE_FIELD,
-            ]
-        ].copy()
-
-        self.data.rename(columns={invoice.SU_HOURS_FIELD: "SU Hours"}, inplace=True)
+        ]

process_report/invoices/nonbillable_invoice.py

@@ -8,6 +8,21 @@ class NonbillableInvoice(invoice.Invoice):
     nonbillable_pis: list[str]
     nonbillable_projects: list[str]
 
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PROJECT_FIELD,
+        invoice.PROJECT_ID_FIELD,
+        invoice.PI_FIELD,
+        invoice.INVOICE_EMAIL_FIELD,
+        invoice.INVOICE_ADDRESS_FIELD,
+        invoice.INSTITUTION_FIELD,
+        invoice.INSTITUTION_ID_FIELD,
+        invoice.SU_HOURS_FIELD,
+        invoice.SU_TYPE_FIELD,
+        invoice.RATE_FIELD,
+        invoice.COST_FIELD,
+    ]
+
     def _prepare_export(self):
         self.data = self.data[
             self.data[invoice.PI_FIELD].isin(self.nonbillable_pis)

process_report/process_report.py

@@ -232,15 +232,15 @@ def main():
     nonbillable_inv = nonbillable_invoice.NonbillableInvoice(
         name=args.nonbillable_file,
         invoice_month=invoice_month,
-        data=add_institute_proc.data,
+        data=lenovo_proc.data,
         nonbillable_pis=pi,
         nonbillable_projects=projects,
     )
 
     ### Perform main processing
 
     remove_nonbillables_proc = remove_nonbillables_processor.RemoveNonbillables(
-        "", invoice_month, add_institute_proc.data, pi, projects
+        "", invoice_month, lenovo_proc.data, pi, projects
     )
     remove_nonbillables_proc.process()
 
@@ -272,19 +272,15 @@ def main():
     billable_inv = billable_invoice.BillableInvoice(
         name=args.output_file,
         invoice_month=invoice_month,
-        data=new_pi_credit_proc.data,
+        data=bu_subsidy_proc.data,
         old_pi_filepath=old_pi_file,
         updated_old_pi_df=new_pi_credit_proc.updated_old_pi_df,
     )
 
-    process_and_export_invoices(
-        [lenovo_inv, nonbillable_inv, billable_inv], args.upload_to_s3
-    )
-
     nerc_total_inv = NERC_total_invoice.NERCTotalInvoice(
         name=args.NERC_total_invoice_file,
         invoice_month=invoice_month,
-        data=new_pi_credit_proc.data.copy(),
+        data=bu_subsidy_proc.data.copy(),
     )
 
     bu_internal_inv = bu_internal_invoice.BUInternalInvoice(
@@ -294,7 +290,10 @@ def main():
         subsidy_amount=args.BU_subsidy_amount,
     )
 
-    process_and_export_invoices([nerc_total_inv, bu_internal_inv], args.upload_to_s3)
+    process_and_export_invoices(
+        [lenovo_inv, nonbillable_inv, billable_inv, nerc_total_inv, bu_internal_inv],
+        args.upload_to_s3,
+    )
 
     export_pi_billables(billable_inv.data.copy(), args.output_folder, invoice_month)
 

process_report/tests/unit_tests.py

@@ -907,3 +907,15 @@ def test_remove_prefix(self, mock_get_time, mock_get_bucket):
         process_report.upload_to_s3(filenames, invoice_month)
         for i, call_args in enumerate(mock_bucket.upload_file.call_args_list):
             self.assertTrue(answers[i] in call_args)
+
+
+class TestBaseInvoice(TestCase):
+    def test_filter_exported_columns(self):
+        test_invoice = pandas.DataFrame(columns=["C1", "C2", "C3", "C4", "C5"])
+        answer_invoice = pandas.DataFrame(columns=["C1", "C3R", "C5R"])
+        inv = test_utils.new_base_invoice(data=test_invoice)
+        inv.export_columns_list = ["C1", "C3", "C5"]
+        inv.exported_columns_map = {"C3": "C3R", "C5": "C5R"}
+        inv._filter_columns()
+
+        self.assertTrue(inv.data.equals(answer_invoice))

process_report/tests/util.py

@@ -9,6 +9,7 @@
     new_pi_credit_processor,
     bu_subsidy_processor,
 )
+from process_report.invoices import invoice
 
 
 def new_add_institution_processor(
@@ -71,3 +72,11 @@ def new_bu_subsidy_processor(
     return bu_subsidy_processor.BUSubsidyProcessor(
         name, invoice_month, data, subsidy_amount
     )
+
+
+def new_base_invoice(
+    name="",
+    invoice_month="0000-00",
+    data=pandas.DataFrame(),
+):
+    return invoice.Invoice(name, invoice_month, data)