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

@@ -22,6 +22,25 @@ class BillableInvoice(discount_invoice.DiscountInvoice):
 
     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,
+        invoice.CREDIT_FIELD,
+        invoice.CREDIT_CODE_FIELD,
+        invoice.BALANCE_FIELD,
+    ]
+
     old_pi_filepath: str
 
     @staticmethod

process_report/invoices/bu_internal_invoice.py

@@ -7,6 +7,16 @@
 
 @dataclass
 class BUInternalInvoice(discount_invoice.DiscountInvoice):
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PI_FIELD,
+        "Project",
+        invoice.COST_FIELD,
+        invoice.CREDIT_FIELD,
+        invoice.SUBSIDY_FIELD,
+        invoice.BALANCE_FIELD,
+    ]
+
     subsidy_amount: int
 
     def _prepare(self):
@@ -22,17 +32,6 @@ def get_project(row):
         ].copy()
         self.data["Project"] = self.data.apply(get_project, axis=1)
         self.data[invoice.SUBSIDY_FIELD] = Decimal(0)
-        self.data = self.data[
-            [
-                invoice.INVOICE_DATE_FIELD,
-                invoice.PI_FIELD,
-                "Project",
-                invoice.COST_FIELD,
-                invoice.CREDIT_FIELD,
-                invoice.SUBSIDY_FIELD,
-                invoice.BALANCE_FIELD,
-            ]
-        ]
 
     def _process(self):
         data_summed_projects = self._sum_project_allocations(self.data)

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"
@@ -33,6 +34,9 @@
 
 @dataclass
 class Invoice:
+    export_columns_list = list()
+    exported_columns_map = dict()
+
     name: str
     invoice_month: str
     data: pandas.DataFrame
@@ -78,7 +82,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

@@ -8,21 +8,24 @@ class LenovoInvoice(invoice.Invoice):
     LENOVO_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
     SU_CHARGE_MULTIPLIER = 1
 
-    def _prepare(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()
+    export_columns_list = [
+        invoice.INVOICE_DATE_FIELD,
+        invoice.PROJECT_FIELD,
+        invoice.INSTITUTION_FIELD,
+        invoice.SU_HOURS_FIELD,
+        invoice.SU_TYPE_FIELD,
+        "SU Charge",
+        "Charge",
+    ]
+    exported_columns_map = {invoice.SU_HOURS_FIELD: "SU Hours"}
 
-        self.data.rename(columns={invoice.SU_HOURS_FIELD: "SU Hours"}, inplace=True)
-        self.data.insert(len(self.data.columns), "SU Charge", self.SU_CHARGE_MULTIPLIER)
+    def _prepare(self):
+        self.data["SU Charge"] = self.SU_CHARGE_MULTIPLIER
 
     def _process(self):
-        self.data["Charge"] = self.data["SU Hours"] * self.data["SU Charge"]
+        self.data["Charge"] = self.data[invoice.SU_HOURS_FIELD] * self.data["SU Charge"]
+
+    def _prepare_export(self):
+        self.data = self.data[
+            self.data[invoice.SU_TYPE_FIELD].isin(self.LENOVO_SU_TYPES)
+        ]

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/tests/unit_tests.py

@@ -773,7 +773,7 @@ def test_process_lenovo(self):
                     process_report.PROJECT_FIELD,
                     process_report.INSTITUTION_FIELD,
                     process_report.SU_TYPE_FIELD,
-                    "SU Hours",
+                    process_report.SU_HOURS_FIELD,
                     "SU Charge",
                     "Charge",
                 ]
@@ -785,7 +785,9 @@ def test_process_lenovo(self):
                 row[process_report.SU_TYPE_FIELD],
                 ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"],
             )
-            self.assertEqual(row["Charge"], row["SU Charge"] * row["SU Hours"])
+            self.assertEqual(
+                row["Charge"], row["SU Charge"] * row["SU Hours (GBhr or SUhr)"]
+            )
 
 
 class TestUploadToS3(TestCase):
@@ -833,3 +835,15 @@ def test_upload_to_s3(self, mock_get_time, mock_get_bucket):
 
         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))