Refactored the HU-BU invoice

process_report/invoices/hu_bu_invoice.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+import process_report.invoices.invoice as invoice
+import process_report.util as util
+
+
+@dataclass
+class HUBUInvoice(invoice.Invoice):
+    @property
+    def output_s3_key(self) -> str:
+        return (
+            f"Invoices/{self.invoice_month}/NERC-{self.invoice_month}-Total-Invoice.csv"
+        )
+
+    @property
+    def output_s3_archive_key(self):
+        return f"Invoices/{self.invoice_month}/Archive/NERC-{self.invoice_month}-Total-Invoice {util.get_iso8601_time()}.csv"
+
+    def _prepare_export(self):
+        self.data = self.data[
+            (self.data[invoice.INSTITUTION_FIELD] == "Harvard University")
+            | (self.data[invoice.INSTITUTION_FIELD] == "Boston University")
+        ].copy()

process_report/process_report.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import datetime
+import functools
 from decimal import Decimal
 
 import json
@@ -13,6 +14,7 @@
     lenovo_invoice,
     nonbillable_invoice,
     billable_invoice,
+    hu_bu_invoice,
 )
 
 
@@ -85,6 +87,7 @@ def load_alias(alias_file):
     return alias_dict
 
 
+@functools.lru_cache
 def get_invoice_bucket():
     try:
         s3_resource = boto3.resource(
@@ -172,7 +175,7 @@ def main():
     parser.add_argument(
         "--HU-BU-invoice-file",
         required=False,
-        default="HU_BU.csv",
+        default="HU_BU",
         help="Name of output csv for HU and BU invoice",
     )
     parser.add_argument(
@@ -266,13 +269,22 @@ def main():
     )
     billable_inv.process()
     billable_inv.export()
+
+    hu_bu_inv = hu_bu_invoice.HUBUInvoice(
+        name=args.HU_BU_invoice_file,
+        invoice_month=invoice_month,
+        data=billable_inv.data,
+    )
+    hu_bu_inv.process()
+    hu_bu_inv.export()
+
     if args.upload_to_s3:
-        bucket = get_invoice_bucket()
-        billable_inv.export_s3(bucket)
+        for invoice in [billable_inv, hu_bu_inv]:
+            bucket = get_invoice_bucket()
+            invoice.export_s3(bucket)
 
     export_pi_billables(billable_inv.data, args.output_folder, invoice_month)
     export_BU_only(billable_inv.data, args.BU_invoice_file, args.BU_subsidy_amount)
-    export_HU_BU(billable_inv.data, args.HU_BU_invoice_file)
 
     if args.upload_to_s3:
         invoice_list = list()
@@ -281,7 +293,6 @@ def main():
             invoice_list.append(os.path.join(args.output_folder, pi_invoice))
 
         upload_to_s3(invoice_list, invoice_month)
-        upload_to_s3_HU_BU(args.HU_BU_invoice_file, invoice_month)
 
 
 def fetch_s3_invoices(invoice_month):
@@ -472,14 +483,6 @@ def _apply_subsidy(dataframe, subsidy_amount):
     return dataframe
 
 
-def export_HU_BU(dataframe, output_file):
-    HU_BU_projects = dataframe[
-        (dataframe[INSTITUTION_FIELD] == "Harvard University")
-        | (dataframe[INSTITUTION_FIELD] == "Boston University")
-    ]
-    HU_BU_projects.to_csv(output_file, index=False)
-
-
 def upload_to_s3(invoice_list: list, invoice_month):
     invoice_bucket = get_invoice_bucket()
     for invoice_filename in invoice_list:
@@ -492,17 +495,5 @@ def upload_to_s3(invoice_list: list, invoice_month):
         invoice_bucket.upload_file(invoice_filename, invoice_s3_path_archive)
 
 
-def upload_to_s3_HU_BU(invoice_filename, invoice_month):
-    invoice_bucket = get_invoice_bucket()
-    invoice_bucket.upload_file(
-        invoice_filename,
-        f"Invoices/{invoice_month}/NERC-{invoice_month}-Total-Invoice.csv",
-    )
-    invoice_bucket.upload_file(
-        invoice_filename,
-        f"Invoices/{invoice_month}/Archive/NERC-{invoice_month}-Total-Invoice {get_iso8601_time()}.csv",
-    )
-
-
 if __name__ == "__main__":
     main()