Implemented discount_invoice to refactor BU-Internal and billable invoice

The BU-Internal and billable invoice now subclasses from `discount_invoice`,
an invoice class which implements a function to apply a flat discount
on a PI's projects. This reduces some code redundancy since the
New-PI credit and the BU subsidy share some similar logic.

Additional smaller changes is done to improve code readability.

Author: QuanMPhm

process_report/invoices/billable_invoice.py

@@ -1,21 +1,20 @@
 from dataclasses import dataclass
-from decimal import Decimal
 import logging
 import sys
 
 import pandas
 import pyarrow
 
-import process_report.invoices.invoice as invoice
-import process_report.util as util
+from process_report.invoices import invoice, discount_invoice
+from process_report import util
 
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
 @dataclass
-class BillableInvoice(invoice.Invoice):
+class BillableInvoice(discount_invoice.DiscountInvoice):
     NEW_PI_CREDIT_CODE = "0002"
     INITIAL_CREDIT_AMOUNT = 1000
     EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
@@ -89,7 +88,7 @@ def _prepare(self):
         self.data = self._validate_pi_names(self.data)
         self.data[invoice.CREDIT_FIELD] = None
         self.data[invoice.CREDIT_CODE_FIELD] = None
-        self.data[invoice.BALANCE_FIELD] = Decimal(0)
+        self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD]
         self.old_pi_df = self._load_old_pis(self.old_pi_filepath)
 
     def _process(self):
@@ -143,14 +142,17 @@ def get_initial_credit_amount(
 
         current_pi_set = set(data[invoice.PI_FIELD])
         for pi in current_pi_set:
-            pi_projects = data[data[invoice.PI_FIELD] == pi]
+            credit_eligible_projects = data[
+                (data[invoice.PI_FIELD] == pi)
+                & ~(data[invoice.SU_TYPE_FIELD].isin(self.EXCLUDE_SU_TYPES))
+            ]
             pi_age = self._get_pi_age(old_pi_df, pi, self.invoice_month)
             pi_old_pi_entry = old_pi_df.loc[
                 old_pi_df[invoice.PI_PI_FIELD] == pi
             ].squeeze()
 
             if pi_age > 1:
-                for i, row in pi_projects.iterrows():
+                for i, row in credit_eligible_projects.iterrows():
                     data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD]
             else:
                 if pi_age == 0:
@@ -176,25 +178,16 @@ def get_initial_credit_amount(
                     )
                     credit_used_field = invoice.PI_2ND_USED
 
-                initial_credit = remaining_credit
-                for i, row in pi_projects.iterrows():
-                    if (
-                        remaining_credit == 0
-                        or row[invoice.SU_TYPE_FIELD] in self.EXCLUDE_SU_TYPES
-                    ):
-                        data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD]
-                    else:
-                        project_cost = row[invoice.COST_FIELD]
-                        applied_credit = min(project_cost, remaining_credit)
-
-                        data.at[i, invoice.CREDIT_FIELD] = applied_credit
-                        data.at[i, invoice.CREDIT_CODE_FIELD] = self.NEW_PI_CREDIT_CODE
-                        data.at[i, invoice.BALANCE_FIELD] = (
-                            row[invoice.COST_FIELD] - applied_credit
-                        )
-                        remaining_credit -= applied_credit
+                credits_used = self.apply_flat_discount(
+                    data,
+                    credit_eligible_projects,
+                    remaining_credit,
+                    invoice.CREDIT_FIELD,
+                    invoice.BALANCE_FIELD,
+                    invoice.CREDIT_CODE_FIELD,
+                    self.NEW_PI_CREDIT_CODE,
+                )
 
-                credits_used = initial_credit - remaining_credit
                 if (pi_old_pi_entry[credit_used_field] != 0) and (
                     credits_used != pi_old_pi_entry[credit_used_field]
                 ):

process_report/invoices/bu_internal_invoice.py

@@ -2,10 +2,11 @@
 from decimal import Decimal
 
 import process_report.invoices.invoice as invoice
+import process_report.invoices.discount_invoice as discount_invoice
 
 
 @dataclass
-class BUInternalInvoice(invoice.Invoice):
+class BUInternalInvoice(discount_invoice.DiscountInvoice):
     subsidy_amount: int
 
     def _prepare(self):
@@ -34,35 +35,36 @@ def get_project(row):
         ]
 
     def _process(self):
-        project_list = self.data["Project"].unique()
-        data_no_dup = self.data.drop_duplicates("Project", inplace=False)
+        data_summed_projects = self._sum_project_allocations(self.data)
+        self.data = self._apply_subsidy(data_summed_projects, self.subsidy_amount)
+
+    def _sum_project_allocations(self, dataframe):
+        """A project may have multiple allocations, and therefore multiple rows
+        in the raw invoices. For BU-Internal invoice, we only want 1 row for
+        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]
         for project in project_list:
-            project_mask = self.data["Project"] == project
+            project_mask = dataframe["Project"] == project
             no_dup_project_mask = data_no_dup["Project"] == project
 
-            sum_fields_sums = self.data[project_mask][sum_fields].sum().values
+            sum_fields_sums = dataframe[project_mask][sum_fields].sum().values
             data_no_dup.loc[no_dup_project_mask, sum_fields] = sum_fields_sums
 
-        self.data = self._apply_subsidy(data_no_dup, self.subsidy_amount)
+        return data_no_dup
 
     def _apply_subsidy(self, dataframe, subsidy_amount):
         pi_list = dataframe[invoice.PI_FIELD].unique()
 
         for pi in pi_list:
             pi_projects = dataframe[dataframe[invoice.PI_FIELD] == pi]
-            remaining_subsidy = subsidy_amount
-            for i, row in pi_projects.iterrows():
-                project_remaining_cost = row[invoice.BALANCE_FIELD]
-                applied_subsidy = min(project_remaining_cost, remaining_subsidy)
-
-                dataframe.at[i, invoice.SUBSIDY_FIELD] = applied_subsidy
-                dataframe.at[i, invoice.BALANCE_FIELD] = (
-                    row[invoice.BALANCE_FIELD] - applied_subsidy
-                )
-                remaining_subsidy -= applied_subsidy
-
-                if remaining_subsidy == 0:
-                    break
+            self.apply_flat_discount(
+                dataframe,
+                pi_projects,
+                subsidy_amount,
+                invoice.SUBSIDY_FIELD,
+                invoice.BALANCE_FIELD,
+            )
 
         return dataframe

process_report/invoices/discount_invoice.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+
+import pandas
+
+import process_report.invoices.invoice as invoice
+
+
+@dataclass
+class DiscountInvoice(invoice.Invoice):
+    """
+    Invoice class containing functions useful for applying discounts
+    on dataframes
+    """
+
+    @staticmethod
+    def apply_flat_discount(
+        invoice: pandas.DataFrame,
+        pi_projects: pandas.DataFrame,
+        discount_amount: int,
+        discount_field: str,
+        balance_field: str,
+        code_field: str = None,
+        discount_code: str = None,
+    ):
+        """
+        Takes in an invoice and a list of PI projects that are a subset of it,
+        and applies a flat discount to those PI projects. Note that this function
+        will change the provided `invoice` Dataframe directly. Therefore, it does
+        not return the changed invoice.
+
+        This function assumes that the balance field shows the remaining cost of the project,
+        or what the PI would pay before the flat discount is applied.
+
+        If the optional parameters `code_field` and `discount_code` are passed in,
+        `discount_code` will be comma-APPENDED to the `code_field` of projects where
+        the discount is applied
+
+        Returns the amount of discount used.
+
+        :param invoice: Dataframe containing all projects
+        :param pi_projects: A subset of `invoice`, containing all projects for a PI you want to apply the discount
+        :param discount_amount: The discount given to the PI
+        :param discount_field: Name of the field to put the discount amount applied to each project
+        :param balance_field: Name of the balance field
+        :param code_field: Name of the discount code field
+        :param discount_code: Code of the discount
+        """
+
+        def apply_discount_on_project(remaining_discount_amount, project_i, project):
+            remaining_project_balance = project[balance_field]
+            applied_discount = min(remaining_project_balance, remaining_discount_amount)
+            invoice.at[project_i, discount_field] = applied_discount
+            invoice.at[project_i, balance_field] = (
+                project[balance_field] - applied_discount
+            )
+            remaining_discount_amount -= applied_discount
+            return remaining_discount_amount
+
+        def apply_credit_code_on_project(project_i):
+            if code_field and discount_code:
+                if pandas.isna(invoice.at[project_i, code_field]):
+                    invoice.at[project_i, code_field] = discount_code
+                else:
+                    invoice.at[project_i, code_field] = (
+                        invoice.at[project_i, code_field] + "," + discount_code
+                    )
+
+        remaining_discount_amount = discount_amount
+        for i, row in pi_projects.iterrows():
+            if remaining_discount_amount == 0:
+                break
+            else:
+                remaining_discount_amount = apply_discount_on_project(
+                    remaining_discount_amount, i, row
+                )
+                apply_credit_code_on_project(i)
+
+        discount_used = discount_amount - remaining_discount_amount
+        return discount_used

process_report/process_report.py

@@ -2,13 +2,12 @@
 import os
 import sys
 import datetime
-import functools
 
 import json
 import pandas
-import boto3
 import pyarrow
 
+from process_report.util import get_invoice_bucket, process_and_export_invoices
 from process_report.invoices import (
     lenovo_invoice,
     nonbillable_invoice,
@@ -87,22 +86,6 @@ def load_alias(alias_file):
     return alias_dict
 
 
-@functools.lru_cache
-def get_invoice_bucket():
-    try:
-        s3_resource = boto3.resource(
-            service_name="s3",
-            endpoint_url=os.environ.get(
-                "S3_ENDPOINT", "https://s3.us-east-005.backblazeb2.com"
-            ),
-            aws_access_key_id=os.environ["S3_KEY_ID"],
-            aws_secret_access_key=os.environ["S3_APP_KEY"],
-        )
-    except KeyError:
-        print("Error: Please set the environment variables S3_KEY_ID and S3_APP_KEY")
-    return s3_resource.Bucket(os.environ.get("S3_BUCKET_NAME", "nerc-invoicing"))
-
-
 def get_iso8601_time():
     return datetime.datetime.now().strftime("%Y%m%dT%H%M%SZ")
 
@@ -249,12 +232,6 @@ def main():
         nonbillable_pis=pi,
         nonbillable_projects=projects,
     )
-    for invoice in [lenovo_inv, nonbillable_inv]:
-        invoice.process()
-        invoice.export()
-        if args.upload_to_s3:
-            bucket = get_invoice_bucket()
-            invoice.export_s3(bucket)
 
     if args.upload_to_s3:
         backup_to_s3_old_pi_file(old_pi_file)
@@ -267,35 +244,27 @@ def main():
         nonbillable_projects=projects,
         old_pi_filepath=old_pi_file,
     )
-    billable_inv.process()
-    billable_inv.export()
+
+    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=billable_inv.data,
+        data=billable_inv.data.copy(),
     )
-    nerc_total_inv.process()
-    nerc_total_inv.export()
-
-    if args.upload_to_s3:
-        for invoice in [billable_inv, nerc_total_inv]:
-            bucket = get_invoice_bucket()
-            invoice.export_s3(bucket)
 
     bu_internal_inv = bu_internal_invoice.BUInternalInvoice(
         name=args.BU_invoice_file,
         invoice_month=invoice_month,
-        data=billable_inv.data,
+        data=billable_inv.data.copy(),
         subsidy_amount=args.BU_subsidy_amount,
     )
-    bu_internal_inv.process()
-    bu_internal_inv.export()
-    if args.upload_to_s3:
-        bucket = get_invoice_bucket()
-        bu_internal_inv.export_s3(bucket)
 
-    export_pi_billables(billable_inv.data, args.output_folder, invoice_month)
+    process_and_export_invoices([nerc_total_inv, bu_internal_inv], args.upload_to_s3)
+
+    export_pi_billables(billable_inv.data.copy(), args.output_folder, invoice_month)
 
     if args.upload_to_s3:
         invoice_list = list()

process_report/tests/unit_tests.py

@@ -5,7 +5,6 @@
 import os
 import uuid
 import math
-from decimal import Decimal
 from textwrap import dedent
 
 from process_report import process_report, util
@@ -419,7 +418,7 @@ def setUp(self):
         self.dataframe = pandas.DataFrame(data)
         self.dataframe["Credit"] = None
         self.dataframe["Credit Code"] = None
-        self.dataframe["Balance"] = Decimal(0)
+        self.dataframe["Balance"] = self.dataframe["Cost"]
         self.answer_dataframe = pandas.DataFrame(answer_df_dict)
         old_pi = [
             "PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used",
@@ -516,7 +515,7 @@ def setUp(self):
         )
         self.dataframe_no_gpu["Credit"] = None
         self.dataframe_no_gpu["Credit Code"] = None
-        self.dataframe_no_gpu["Balance"] = Decimal(0)
+        self.dataframe_no_gpu["Balance"] = self.dataframe_no_gpu["Cost"]
         old_pi_no_gpu = [
             "PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used",
             "OldPI,2024-03,500,200,0",