update import volunteer script to be more flexible

crowdsourcer/management/commands/import_volunteers.py

@@ -7,7 +7,14 @@
 
 import pandas as pd
 
-from crowdsourcer.models import Assigned, PublicAuthority, Section
+from crowdsourcer.models import (
+    Assigned,
+    Marker,
+    MarkingSession,
+    PublicAuthority,
+    ResponseType,
+    Section,
+)
 
 YELLOW = "\033[33m"
 RED = "\033[31m"
@@ -18,10 +25,12 @@
 class Command(BaseCommand):
     help = "import volunteers"
 
-    volunteer_file = settings.BASE_DIR / "data" / "volunteers.xlsx"
+    volunteer_file = settings.BASE_DIR / "data" / "volunteers.csv"
+    response_type = "First Mark"
 
     section_map = {
         "Buildings": "Buildings & Heating",
+        "Buildings and Heating": "Buildings & Heating",
         "Planning": "Planning & Land Use",
         "Gov & Finance": "Governance & Finance",
         "Waste & Food": "Waste Reduction & Food",
@@ -30,6 +39,7 @@ class Command(BaseCommand):
 
     num_council_map = {
         "scorecards_volunteering": 6,
+        "scorecards_assessor": 6,
         "local_climate_policy_programme": 15,
     }
 
@@ -42,11 +52,46 @@ class Command(BaseCommand):
         "assigned_section",
     ]
 
+    usecols = [
+        "First name",
+        "Last name",
+        "Email",
+        "Council Area",
+        "Type of Volunteering",
+        "Assigned Section",
+    ]
+
     def add_arguments(self, parser):
         parser.add_argument(
             "-q", "--quiet", action="store_true", help="Silence progress bars."
         )
 
+        parser.add_argument(
+            "--file",
+            action="store",
+            required=True,
+            help="CSV file containing the assignments",
+        )
+
+        parser.add_argument(
+            "--col_names",
+            action="store",
+            help="CSV file containing use_cols and col_names args (one column for each)",
+        )
+
+        parser.add_argument(
+            "--response_type",
+            action="store",
+            help="Stage to assign markers to",
+        )
+
+        parser.add_argument(
+            "--session",
+            action="store",
+            required=True,
+            help="Marking session to use assignements with",
+        )
+
         parser.add_argument(
             "--add_users", action="store_true", help="add users to database"
         )
@@ -55,21 +100,51 @@ def add_arguments(self, parser):
             "--make_assignments", action="store_true", help="assign councils to users"
         )
 
-    def handle(self, quiet: bool = False, *args, **options):
-        df = pd.read_excel(
-            self.volunteer_file,
-            usecols=[
-                "First name",
-                "Last name",
-                "Email",
-                "Council Area",
-                "Type of Volunteering",
-                "Assigned Section",
-            ],
-            sheet_name="Volunteer Recruitment Cohort 2",
+    def get_df(self, filename):
+        df = pd.read_csv(
+            filename,
+            usecols=self.usecols,
         )
         df.columns = self.column_names
 
+        return df
+
+    def set_cols(self, col_names):
+        if col_names is not None:
+            cols = pd.read_csv(col_names)
+            self.usecols = cols.use_cols
+            self.column_names = cols.col_names
+
+    def get_assignment_count(self, user_type):
+        user_type = user_type.lower().replace(" ", "_")
+
+        num_councils = self.num_council_map.get(user_type)
+        return num_councils
+
+    def handle(
+        self,
+        quiet: bool = False,
+        file: str = None,
+        session: str = None,
+        col_names: str = None,
+        response_type: str = None,
+        *args,
+        **options,
+    ):
+        if file is None:
+            file = self.volunteer_file
+
+        self.set_cols(col_names)
+
+        df = self.get_df(file)
+
+        if response_type is None:
+            response_type = self.response_type
+
+        session = MarkingSession.objects.get(label=session)
+        rt = ResponseType.objects.get(type=response_type)
+
+        bad_councils = []
         for index, row in df.iterrows():
             if pd.isna(row["email"]):
                 continue
@@ -82,8 +157,15 @@ def handle(self, quiet: bool = False, *args, **options):
                 self.stdout.write(f"{YELLOW}No user type for {row['email']}{NOBOLD}")
                 continue
 
+            num_councils = self.get_assignment_count(user_type)
+            if num_councils is None:
+                self.stdout.write(
+                    f"{YELLOW}Don't know how many councils to assign for {row['email']}{NOBOLD}"
+                )
+                continue
+
             if options["add_users"] is True:
-                u, created = User.objects.update_or_create(
+                u, _ = User.objects.update_or_create(
                     username=row["email"],
                     defaults={
                         "email": row["email"],
@@ -92,6 +174,13 @@ def handle(self, quiet: bool = False, *args, **options):
                     },
                 )
                 u.save()
+                m, _ = Marker.objects.update_or_create(
+                    user=u,
+                    defaults={
+                        "response_type": rt,
+                    },
+                )
+                m.marking_session.set([session])
             else:
                 try:
                     u = User.objects.get(username=row["email"])
@@ -106,7 +195,7 @@ def handle(self, quiet: bool = False, *args, **options):
                     title = self.section_map.get(
                         row["assigned_section"], row["assigned_section"]
                     )
-                    s = Section.objects.get(title=title)
+                    s = Section.objects.get(title=title, marking_session=session)
                 except Section.DoesNotExist:
                     self.stdout.write(
                         f"{RED}could not assign section for {row['email']}, no section {title}{NOBOLD}"
@@ -134,6 +223,7 @@ def handle(self, quiet: bool = False, *args, **options):
                     )
 
             if len(councils) > 0 and own_council.count() == 0:
+                bad_councils.append((row["council_area"], row["email"]))
                 self.stdout.write(
                     f"{RED}Bad council: {row['council_area']} (f{row['email']}){NOBOLD}"
                 )
@@ -148,11 +238,13 @@ def handle(self, quiet: bool = False, *args, **options):
             own_council_list = list(own_council.values_list("id", flat=True))
             assigned_councils = assigned_councils + own_council_list
 
-            num_councils = self.num_council_map[user_type]
-
-            councils_to_assign = PublicAuthority.objects.exclude(
+            councils_to_assign = PublicAuthority.objects.filter(
+                marking_session=session
+            ).exclude(
                 Q(id__in=assigned_councils) | Q(type="COMB") | Q(do_not_mark=True)
-            )[:num_councils]
+            )[
+                :num_councils
+            ]
 
             if councils_to_assign.count() == 0:
                 self.stdout.write(
@@ -162,11 +254,13 @@ def handle(self, quiet: bool = False, *args, **options):
             if options["make_assignments"] is True:
                 for council in councils_to_assign:
                     a, created = Assigned.objects.update_or_create(
-                        user=u, section=s, authority=council
+                        user=u, section=s, authority=council, marking_session=session
                     )
 
-        council_count = PublicAuthority.objects.filter(do_not_mark=False).count()
-        for section in Section.objects.all():
+        council_count = PublicAuthority.objects.filter(
+            marking_session=session, do_not_mark=False
+        ).count()
+        for section in Section.objects.filter(marking_session=session).all():
             assigned = Assigned.objects.filter(section=section).count()
             if assigned != council_count:
                 self.stdout.write(
@@ -191,3 +285,8 @@ def handle(self, quiet: bool = False, *args, **options):
             self.stdout.write(
                 f"{YELLOW}Dry run, no assignments made, call with --make_assignments to make them{NOBOLD}"
             )
+
+        if len(bad_councils):
+            self.stdout.write("Bad councils are:")
+            for c in bad_councils:
+                self.stdout.write(f"{YELLOW}{c[0]}, {c[1]}{NOBOLD}")

crowdsourcer/tests/data/volunteers.csv

@@ -0,0 +1,2 @@
+First name,Last name,Email,Council Area,Type of Volunteering,Assigned Section
+First,Last,[email protected],Aberdeenshire Council,Scorecards Volunteering,Transport

crowdsourcer/tests/data/volunteers_multiple.csv

@@ -0,0 +1,3 @@
+First name,Last name,Email,Council Area,Type of Volunteering,Assigned Section
+First,Last,[email protected],Aberdeenshire Council,Scorecards Volunteering,Transport
+Primary,Secondary,[email protected],Aberdeen City Council,Scorecards Volunteering,Transport