Charging structure changes

Saket's charging structure changes made into a merge request, rather than direct commits on develop

compendium_v2/publishers/survey_publisher_2022.py

@@ -5,6 +5,8 @@ import math
 
 from sqlalchemy import text
 from collections import defaultdict
+
+from compendium_v2.db.model import FeeType
 from compendium_v2.environment import setup_logging
 from compendium_v2.config import load
 from compendium_v2 import db, survey_db
@@ -92,6 +94,13 @@ class OrgQuestion(enum.Enum):
     sub_orgs_5_role = 16439
 
 
+class ChargingStructure(enum.Enum):
+    """
+    Answers are strings
+    """
+    charging_structure = 16410
+
+
 def query_budget():
     with survey_db.session_scope() as survey:
         return survey.execute(text(BUDGET_QUERY))
@@ -214,23 +223,28 @@ def transfer_staff_data():
                     continue
 
                 # initialize on first use, so we don't add data for nrens with no answers
-                data.setdefault(nren_name, {question: 0 for question in StaffQuestion})[question] = value
+                data.setdefault(nren_name,
+                                {question: 0 for question in StaffQuestion})[
+                    question] = value
 
         for nren_name, nren_info in data.items():
             if sum([nren_info[question] for question in StaffQuestion]) == 0:
-                logger.info(f'{nren_name} has no staff data. Deleting if exists.')
+                logger.info(
+                    f'{nren_name} has no staff data. Deleting if exists.')
                 session.query(model.NrenStaff).filter(
                     model.NrenStaff.nren_id == nren_dict[nren_name].id,
                     model.NrenStaff.year == 2022,
                 ).delete()
                 continue
 
-            employed = nren_info[StaffQuestion.permanent_fte] + nren_info[StaffQuestion.subcontracted_fte]
-            technical = nren_info[StaffQuestion.technical_fte] + nren_info[StaffQuestion.non_technical_fte]
+            employed = nren_info[StaffQuestion.permanent_fte] + nren_info[
+                StaffQuestion.subcontracted_fte]
+            technical = nren_info[StaffQuestion.technical_fte] + nren_info[
+                StaffQuestion.non_technical_fte]
             if not math.isclose(
-                employed,
-                technical,
-                abs_tol=0.01,
+                    employed,
+                    technical,
+                    abs_tol=0.01,
             ):
                 logger.info(
                     f'{nren_name} FTE do not equal across employed/technical categories.'
@@ -249,7 +263,6 @@ def transfer_staff_data():
 
 
 def transfer_nren_parent_org():
-
     # clean up the data a bit by removing some strings
     strings_to_replace = [
         'We are affiliated to '
@@ -280,13 +293,17 @@ def transfer_nren_parent_org():
 
 
 def transfer_nren_sub_org():
-
     suborg_questions = [
-        (OrgQuestion.sub_orgs_1_name, OrgQuestion.sub_orgs_1_choice, OrgQuestion.sub_orgs_1_role),
-        (OrgQuestion.sub_orgs_2_name, OrgQuestion.sub_orgs_2_choice, OrgQuestion.sub_orgs_2_role),
-        (OrgQuestion.sub_orgs_3_name, OrgQuestion.sub_orgs_3_choice, OrgQuestion.sub_orgs_3_role),
-        (OrgQuestion.sub_orgs_4_name, OrgQuestion.sub_orgs_4_choice, OrgQuestion.sub_orgs_4_role),
-        (OrgQuestion.sub_orgs_5_name, OrgQuestion.sub_orgs_5_choice, OrgQuestion.sub_orgs_5_role)
+        (OrgQuestion.sub_orgs_1_name, OrgQuestion.sub_orgs_1_choice,
+         OrgQuestion.sub_orgs_1_role),
+        (OrgQuestion.sub_orgs_2_name, OrgQuestion.sub_orgs_2_choice,
+         OrgQuestion.sub_orgs_2_role),
+        (OrgQuestion.sub_orgs_3_name, OrgQuestion.sub_orgs_3_choice,
+         OrgQuestion.sub_orgs_3_role),
+        (OrgQuestion.sub_orgs_4_name, OrgQuestion.sub_orgs_4_choice,
+         OrgQuestion.sub_orgs_4_role),
+        (OrgQuestion.sub_orgs_5_name, OrgQuestion.sub_orgs_5_choice,
+         OrgQuestion.sub_orgs_5_role)
     ]
 
     with db.session_scope() as session:
@@ -329,6 +346,42 @@ def transfer_nren_sub_org():
         session.commit()
 
 
+def transfer_charging_structure():
+    with db.session_scope() as session:
+        nren_dict = helpers.get_uppercase_nren_dict(session)
+
+        rows = query_question(ChargingStructure.charging_structure)
+        for row in rows:
+            nren_name = row[0].upper()
+            value = row[1].replace('"', '').strip()
+
+            if nren_name not in nren_dict:
+                logger.info(f'{nren_name} unknown. Skipping from charging '
+                            f'structure.')
+                continue
+
+            if "do not charge" in value:
+                charging_structure = FeeType.no_charge.value
+            elif "combination" in value:
+                charging_structure = FeeType.combination.value
+            elif "flat" in value:
+                charging_structure = FeeType.flat_fee.value
+            elif "usage-based" in value:
+                charging_structure = FeeType.usage_based_fee.value
+            elif "Other" in value:
+                charging_structure = FeeType.other.value
+            else:
+                charging_structure = None
+
+            charging_structure = model.ChargingStructure(
+                nren_id=nren_dict[nren_name].id,
+                year=2022,
+                fee_type=charging_structure,
+            )
+            session.merge(charging_structure)
+        session.commit()
+
+
 def _cli(config):
     helpers.init_db(config)
     transfer_budget()
@@ -336,6 +389,7 @@ def _cli(config):
     transfer_staff_data()
     transfer_nren_parent_org()
     transfer_nren_sub_org()
+    transfer_charging_structure()
 
 
 @click.command()

[Bjarke Madsen]

@saket.agrahari please change your environment to use pep8 120 lines formatting instead of 79