enable reading section counts from the import data

Author: struan

crowdsourcer/management/commands/import_volunteers.py

@@ -23,6 +23,24 @@
 GREEN = "\033[32m"
 NOBOLD = "\033[0m"
 
+COLUMN_NAMES = (
+    "first_name",
+    "last_name",
+    "email",
+    "council_area",
+    "type_of_volunteering",
+    "assigned_section",
+)
+
+USECOLS = (
+    "First name",
+    "Last name",
+    "Email",
+    "Council Area",
+    "Type of Volunteering",
+    "Assigned Section",
+)
+
 
 class Command(BaseCommand):
     help = "import volunteers"
@@ -45,24 +63,6 @@ class Command(BaseCommand):
         "local_climate_policy_programme": 15,
     }
 
-    column_names = [
-        "first_name",
-        "last_name",
-        "email",
-        "council_area",
-        "type_of_volunteering",
-        "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."
@@ -106,6 +106,12 @@ def add_arguments(self, parser):
             help="Marking session to use assignements with",
         )
 
+        parser.add_argument(
+            "--section_count_in_data",
+            action="store_true",
+            help="use the section count in the csv and not the map",
+        )
+
         parser.add_argument(
             "--add_users", action="store_true", help="add users to database"
         )
@@ -120,6 +126,17 @@ def add_arguments(self, parser):
             help="run everything and then undo changes. Helpful to check if assigment weighting is good",
         )
 
+    def init(self, file, options):
+        self.column_names = list(COLUMN_NAMES)
+        self.usecols = list(USECOLS)
+
+        if file is None:
+            file = self.volunteer_file
+
+        if options["section_count_in_data"]:
+            self.column_names.append("section_count")
+            self.usecols.append("How many sections?")
+
     def get_df(self, filename):
         df = pd.read_csv(
             filename,
@@ -150,7 +167,10 @@ def set_authority_map(self, authority_map_file):
             for _, row in cols.iterrows():
                 self.authority_map[row.bad_name] = row.good_name
 
-    def get_assignment_count(self, user_type):
+    def get_assignment_count(self, user_type, row, section_count_in_data):
+        if section_count_in_data:
+            return row["section_count"]
+
         user_type = user_type.lower().replace(" ", "_")
 
         num_councils = self.num_council_map.get(user_type)
@@ -171,7 +191,9 @@ def add_users_and_assignments(self, df, response_type, session, rt, options):
                 self.stdout.write(f"{YELLOW}No user type for {row['email']}{NOBOLD}")
                 continue
 
-            num_councils = self.get_assignment_count(user_type)
+            num_councils = self.get_assignment_count(
+                user_type, row, options["section_count_in_data"]
+            )
             if num_councils is None:
                 self.stdout.write(
                     f"{YELLOW}Don't know how many councils to assign for {row['email']}{NOBOLD}"
@@ -296,9 +318,7 @@ def handle(
         *args,
         **options,
     ):
-        if file is None:
-            file = self.volunteer_file
-
+        self.init(file, options)
         self.set_cols(col_names)
         self.set_assignment_map(assignment_map)
         self.set_authority_map(authority_map)

crowdsourcer/tests/data/volunteers_with_limit.csv

@@ -0,0 +1,2 @@
+First name,Last name,Email,Council Area,Type of Volunteering,Assigned Section,How many sections?
+First,Last,first_last@example.org,Aberdeenshire Council,Scorecards Volunteering,Transport,3

crowdsourcer/tests/test_import_volunteers_script.py

@@ -414,3 +414,30 @@ def test_council_map(self):
         self.assertTrue("Aberdeenshire Council" not in councils)
         self.assertTrue("Adur District Council" not in councils)
         self.assertTrue("Aberdeen City Council" in councils)
+
+    def test_assignment_limits_in_file(self):
+        data_file = (
+            pathlib.Path(__file__).parent.resolve()
+            / "data"
+            / "volunteers_with_limit.csv"
+        )
+        self.add_extra_councils()
+
+        self.assertEquals(Assigned.objects.count(), 0)
+        self.call_command(
+            "import_volunteers",
+            session="Default",
+            file=data_file,
+            add_users=True,
+            make_assignments=True,
+            section_count_in_data=True,
+        )
+        self.assertEquals(User.objects.count(), 1)
+        self.assertEquals(Marker.objects.count(), 1)
+        self.assertEquals(Assigned.objects.count(), 3)
+        self.assertEquals(
+            Assigned.objects.filter(
+                marking_session__label="Default", section__title="Transport"
+            ).count(),
+            3,
+        )