Refactor institute_list model

- Moved functions related to `InstituteList` into the class

Author: QuanMPhm

process_report/institute_list_models.py

@@ -1,11 +1,16 @@
 from typing import Annotated
 import datetime
 import functools
+import logging
 
 import pydantic
 import validators
 
 
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
 def parse_date(v: str) -> str:
     try:
         datetime.datetime.strptime(v, "%Y-%m")
@@ -70,3 +75,27 @@ def nonbillable_course_list(self) -> list[str]:
             if institute_info.courses_nonbillable:
                 institute_list.append(institute_info.display_name)
         return institute_list
+
+    @functools.cached_property
+    def domain_institute_mapping(self) -> dict[str, str]:
+        """Dict mapping web domains to institution display names"""
+        institute_map = dict()
+        for institute_info in self.root:
+            for domain in institute_info.domains:
+                institute_map[domain] = institute_info.display_name
+
+        return institute_map
+
+    def get_institution_from_pi(self, pi_email) -> str:
+        institution_domain = pi_email.split("@")[-1]
+        for i in range(institution_domain.count(".") + 1):
+            if institution_name := self.domain_institute_mapping.get(
+                institution_domain, ""
+            ):
+                break
+            institution_domain = institution_domain[institution_domain.find(".") + 1 :]
+
+        if institution_name == "":
+            logger.warning(f"PI name {pi_email} does not match any institution!")
+
+        return institution_name

process_report/processors/add_institution_processor.py

@@ -27,15 +27,14 @@ def _add_institution(self):
         The list of mappings are defined in `institute_map.json`.
         """
         institute_list = util.load_institute_list()
-        institute_map = util.get_institute_mapping(institute_list)
         self.data = self.data.astype({invoice.INSTITUTION_FIELD: "str"})
         for i, row in self.data.iterrows():
             pi_name = row[invoice.PI_FIELD]
             if pandas.isna(pi_name):
                 logger.info(f"Project {row[invoice.PROJECT_FIELD]} has no PI")
             else:
                 self.data.at[i, invoice.INSTITUTION_FIELD] = (
-                    util.get_institution_from_pi(institute_map, pi_name)
+                    institute_list.get_institution_from_pi(pi_name)
                 )
 
     def _process(self):

process_report/processors/prepayment_processor.py

@@ -150,11 +150,10 @@ def _get_prepay_group_dict(self):
     def _add_prepay_info(self):
         """Populate prepaid group name, institute, and MGHPCC managed field"""
         institute_list = util.load_institute_list()
-        institute_map = util.get_institute_mapping(institute_list)
 
         for group_name, group_dict in self.group_info_dict.items():
-            group_institute = util.get_institution_from_pi(
-                institute_map, group_dict[invoice.PREPAY_GROUP_CONTACT_FIELD]
+            group_institute = institute_list.get_institution_from_pi(
+                group_dict[invoice.PREPAY_GROUP_CONTACT_FIELD]
             )
 
             # Prepay projects are identified by project name, not project - allocation name

process_report/tests/unit/processors/test_add_institution_processor.py

@@ -2,6 +2,7 @@
 import pytest
 
 from process_report.institute_list_validate import main
+from process_report.institute_list_models import InstituteList
 from process_report.tests.base import BaseTestCaseWithTempDir
 
 
@@ -42,3 +43,36 @@ def test_invalid_institute_list(self):
             f.flush()
             with pytest.raises(SystemExit):
                 main(["--github", str(test_file)])
+
+    def test_get_pi_institution(self):
+        domain_map = {
+            "harvard.edu": "Harvard University",
+            "bu.edu": "Boston University",
+            "bentley.edu": "Bentley",
+            "mclean.harvard.edu": "McLean Hospital",
+            "northeastern.edu": "Northeastern University",
+            "childrens.harvard.edu": "Boston Children's Hospital",
+            "meei.harvard.edu": "Massachusetts Eye & Ear",
+            "dfci.harvard.edu": "Dana-Farber Cancer Institute",
+            "bwh.harvard.edu": "Brigham and Women's Hospital",
+            "bidmc.harvard.edu": "Beth Israel Deaconess Medical Center",
+        }
+        test_institute_list = InstituteList([])
+        test_institute_list.domain_institute_mapping = domain_map
+
+        answers = {
+            "[email protected]": "Boston University",
+            "[email protected]": "McLean Hospital",
+            "[email protected]": "Harvard University",
+            "e@edu": "",
+            "[email protected]": "Northeastern University",
+            "[email protected]": "Harvard University",
+            "[email protected]": "Boston Children's Hospital",
+            "[email protected]": "Massachusetts Eye & Ear",
+            "[email protected]": "",
+            "[email protected]": "Brigham and Women's Hospital",
+            "[email protected]": "Beth Israel Deaconess Medical Center",
+        }
+
+        for pi_email, answer in answers.items():
+            assert test_institute_list.get_institution_from_pi(pi_email) == answer

process_report/tests/unit/test_institute_list_validate.py

@@ -1,16 +1,12 @@
 import os
 import datetime
 import yaml
-import logging
 import functools
 
 import boto3
 
 from process_report.institute_list_models import InstituteList
 
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
 
 DEFAULT_INSTITUTE_LIST = "process_report/institute_list.yaml"
 
@@ -39,28 +35,6 @@ def load_institute_list() -> InstituteList:
         return InstituteList.model_validate(institute_list)
 
 
-def get_institute_mapping(institute_list: InstituteList):
-    institute_map = dict()
-    for institute_info in institute_list.root:
-        for domain in institute_info.domains:
-            institute_map[domain] = institute_info.display_name
-
-    return institute_map
-
-
-def get_institution_from_pi(institute_map, pi_uname):
-    institution_domain = pi_uname.split("@")[-1]
-    for i in range(institution_domain.count(".") + 1):
-        if institution_name := institute_map.get(institution_domain, ""):
-            break
-        institution_domain = institution_domain[institution_domain.find(".") + 1 :]
-
-    if institution_name == "":
-        logger.warning(f"PI name {pi_uname} does not match any institution!")
-
-    return institution_name
-
-
 def get_iso8601_time():
     return datetime.datetime.now().strftime("%Y%m%dT%H%M%SZ")