fixup! Adds division_io.py

scripts/division_io.py

@@ -8,7 +8,7 @@
 """
 
 import re
-from enum import IntEnum
+from enum import Enum
 from pathlib import Path
 from typing import Dict, cast
 from warnings import filterwarnings
@@ -27,10 +27,20 @@
 )
 
 
-class TitlePriority(IntEnum):
-    ORIGINAL_HEADER = 1
-    PARLIAMENT_DESCRIBED = 4
-    MANUAL = 5
+class TitlePriority(str, Enum):
+    ORIGINAL_HEADER = "ORIGINAL_HEADER"
+    PARLIAMENT_DESCRIBED = "PARLIAMENT_DESCRIBED"
+    MANUAL = "MANUAL"
+
+    @classmethod
+    def get_priority(cls, priority: str) -> int:
+
+        lookup = {
+            cls.ORIGINAL_HEADER: 1,
+            cls.PARLIAMENT_DESCRIBED: 5,
+            cls.MANUAL: 10,
+        }
+        return lookup[priority]
 
 
 @click.group()
@@ -104,9 +114,15 @@ def df_to_db(df: pd.DataFrame, *, new_priority: TitlePriority, verbose: bool = F
 
     # get all divisions with a title_priority below or equal to current priority
     existing_df = pd.read_sql(
-        f"SELECT division_id FROM divisions WHERE title_priority < {new_priority}",
+        f"SELECT division_id, title_priority FROM divisions",
         db_connection,
     )
+    existing_df["int_title_priority"] = existing_df["title_priority"].apply(
+        TitlePriority.get_priority
+    )
+    existing_df = existing_df[
+        existing_df["int_title_priority"] <= TitlePriority.get_priority(new_priority)
+    ]
 
     # limit dataframe to just those that can be updated
 
@@ -117,30 +133,28 @@ def df_to_db(df: pd.DataFrame, *, new_priority: TitlePriority, verbose: bool = F
         print(
             f"[blue]Updating titles for {len(df)} divisions - ignoring {count_removed} due to priority or absence."
         )
-
     if optional_col_count == 0:
         # update the division_title column in the database
-        update_command = (
-            "UPDATE divisions SET division_title = %s WHERE division_id = %s"
-        )
+        update_command = "UPDATE divisions SET division_title = %s, title_priority = %s WHERE division_id = %s"
 
         with db_connection.cursor() as cursor:
             update_data = [
-                (row["division_title"], row["division_id"]) for _, row in df.iterrows()
+                (row["division_title"], new_priority, row["division_id"])
+                for _, row in df.iterrows()
             ]
             cursor.executemany(update_command, update_data)
         db_connection.commit()
         if verbose:
             print(f"[green]{len(df)} rows updated.")
-
     elif optional_col_count == 2:
         # update the division_title, yes_text, and no_text columns in the database
-        update_command = "UPDATE divisions SET division_title = %s, yes_text = %s, no_text = %s WHERE division_id = %s"
+        update_command = "UPDATE divisions SET division_title = %s, title_priority = %s, yes_text = %s, no_text = %s WHERE division_id = %s"
 
         with db_connection.cursor() as cursor:
             update_data = [
                 (
                     row["division_title"],
+                    new_priority,
                     row["yes_text"],
                     row["no_text"],
                     row["division_id"],
@@ -257,6 +271,10 @@ def update_division_title():
         ),
     }
 
+    if item["yes_text"] == "":
+        del item["yes_text"]
+    if item["no_text"] == "":
+        del item["no_text"]
     df = pd.DataFrame([item])
 
     df_to_db(df, new_priority=TitlePriority.MANUAL)