Extraction code validates, but tests fail

src/electos/ballotmaker/ballots/demo_ballot.py → src/electos/ballotmaker/ballots/ballot_layout.py

@@ -1,20 +1,19 @@
 """
-The Demo Ballot module contains the document specifications,
+The Ballot Layout module contains the document specifications,
 page templates, and specific pages
 """
+import logging
 from datetime import datetime
 from functools import partial
 from pathlib import Path
 
 from electos.ballotmaker.ballots.contest_layout import (
-    BallotMeasureData,
     BallotMeasureLayout,
-    CandidateContestData,
     CandidateContestLayout,
 )
 from electos.ballotmaker.ballots.instructions import Instructions
 from electos.ballotmaker.ballots.page_layout import PageLayout
-from electos.ballotmaker.demo_data import spacetown_data
+from electos.ballotmaker.data.models import BallotStyleData
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.units import inch
 from reportlab.platypus import (
@@ -27,6 +26,14 @@
 )
 from reportlab.platypus.flowables import CondPageBreak
 
+logging.getLogger(__name__)
+
+# TODO: use enums in ContestType: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973606562
+# Also: see line 147
+CANDIDATE = "candidate"
+BALLOT_MEASURE = "ballot measure"
+
+
 # set up frames
 # 1 = True, 0 = FALSE
 SHOW_BOUNDARY = 0
@@ -71,31 +78,22 @@
 )
 
 
-def get_election_header() -> dict:
-    return {
-        "Name": "General Election",
-        "StartDate": "2024-11-05",
-        "EndDate": "2024-11-05",
-        "Type": "general",
-        "ElectionScope": "Spacetown Precinct",
-    }
-
-
-def add_header_line(font_size, line_text, new_line=False):
+def add_header_line(
+    font_size: int, line_text: str, new_line: bool = False
+) -> str:
     line_end = "<br />" if new_line else ""
     return f"<font size={font_size}><b>{line_text}</b></font>{line_end}"
 
 
-def build_header_text():
-    elect_dict = get_election_header()
+def build_header_text(election_header: dict, scope: str) -> str:
     font_size = 12
     formatted_header = add_header_line(
-        font_size, f"Sample Ballot for {elect_dict['Name']}", new_line=True
+        font_size,
+        f"Sample Ballot for {election_header['Name']}",
+        new_line=True,
     )
-    formatted_header += add_header_line(
-        font_size, elect_dict["ElectionScope"], new_line=True
-    )
-    end_date = datetime.fromisoformat(elect_dict["EndDate"])
+    formatted_header += add_header_line(font_size, scope, new_line=True)
+    end_date = datetime.fromisoformat(election_header["EndDate"])
     formatted_date = end_date.strftime("%B %m, %Y")
     formatted_header += add_header_line(font_size, formatted_date)
 
@@ -110,20 +108,27 @@ def header(canvas, doc, content):
     canvas.restoreState()
 
 
-def build_ballot() -> str:
-    # create PDF filename; include
-    # datestamp string for PDF
+def build_ballot(
+    ballot_data: BallotStyleData, election_header: dict, output_dir: Path
+) -> str:
+    # create PDF filename
     now = datetime.now()
     date_time = now.strftime("%Y_%m_%dT%H%M%S")
-    home_dir = Path.home()
-    ballot_name = f"{home_dir}/ballot_demo_{date_time}.pdf"
+    ballot_label = ballot_data.id
+    ballot_scope_count = len(ballot_data.scopes)
+    if ballot_scope_count > 1:
+        raise NotImplementedError(
+            f"Multiple ballot scopes currently unsupported. Found {ballot_scope_count} ballot scopes."
+        )
+    ballot_scope = ballot_data.scopes[0]
+    ballot_name = f"{output_dir}/{ballot_label}_{date_time}.pdf"
 
     doc = BaseDocTemplate(ballot_name)
 
     styles = getSampleStyleSheet()
     normal = styles["Normal"]
-    head_text = build_header_text()
-    header_content = Paragraph(head_text, normal)
+    header_text = build_header_text(election_header, ballot_scope)
+    header_content = Paragraph(header_text, normal)
     three_column_template = PageTemplate(
         id="3col",
         frames=[left_frame, mid_frame, right_frame],
@@ -139,42 +144,48 @@ def build_ballot() -> str:
     )
     doc.addPageTemplates(three_column_template)
     doc.addPageTemplates(one_column_template)
-    # add a ballot contest to the second frame (colomn)
-    layout_1 = CandidateContestLayout(
-        CandidateContestData(spacetown_data.can_con_1)
-    )
-    layout_2 = CandidateContestLayout(
-        CandidateContestData(spacetown_data.can_con_2)
-    )
-    layout_3 = CandidateContestLayout(
-        CandidateContestData(spacetown_data.can_con_3)
-    )
-    layout_4 = CandidateContestLayout(
-        CandidateContestData(spacetown_data.can_con_4)
-    )
-    layout_5 = BallotMeasureLayout(
-        BallotMeasureData(spacetown_data.ballot_measure_1)
-    )
-    layout_6 = BallotMeasureLayout(
-        BallotMeasureData(spacetown_data.ballot_measure_2)
-    )
+
+    # TODO: use ballot_data.candidate_contests & .ballot_measures instead.
+    # See thread for details: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973608016
+    candidate_contests = []
+    ballot_measures = []
+    # get contests
+    for count, contest in enumerate(ballot_data.contests, start=1):
+        title = contest.title
+        con_type = contest.type
+        logging.info(f"Found contest: {title} - {con_type}")
+        if con_type == CANDIDATE:
+            candidate_contests.append(contest)
+        elif con_type == BALLOT_MEASURE:
+            ballot_measures.append(contest)
+        else:
+            raise ValueError(f"Unknown contest type: {con_type}")
+    logging.info(f"Total: {count} contests.")
+
     elements = []
     # add voting instructions
     inst = Instructions()
-    elements = inst.instruction_list
     elements.append(NextPageTemplate("3col"))
-    elements.append(layout_1.contest_table)
-    elements.append(layout_2.contest_table)
-    elements.append(CondPageBreak(c_height * inch))
-    elements.append(layout_4.contest_table)
-    elements.append(layout_3.contest_table)
+    elements = inst.instruction_list
+
+    # add candidate contests
+    for can_con_count, candidate_contest in enumerate(
+        candidate_contests, start=1
+    ):
+        candidate_layout = CandidateContestLayout(
+            candidate_contest
+        ).contest_table
+        elements.append(candidate_layout)
+        # insert column break after every 2 contests
+        if (can_con_count % 2 == 0) and (can_con_count < 4):
+            elements.append(CondPageBreak(c_height * inch))
+    # TODO: write more informative log message, see: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973608278
+    logging.info(f"Added {can_con_count} candidate contests.")
     elements.append(NextPageTemplate("1col"))
     elements.append(PageBreak())
-    elements.append(layout_5.contest_table)
-    elements.append(layout_6.contest_table)
+    for measures, ballot_measure in enumerate(ballot_measures, start=1):
+        ballot_layout = BallotMeasureLayout(ballot_measure).contest_table
+        elements.append(ballot_layout)
+    logging.info(f"Added {measures} ballot measures.")
     doc.build(elements)
     return str(ballot_name)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    build_ballot()

src/electos/ballotmaker/ballots/build_ballots.py

@@ -0,0 +1,46 @@
+import logging
+from datetime import datetime
+from pathlib import Path
+
+from electos.ballotmaker.ballots.ballot_layout import build_ballot
+from electos.ballotmaker.constants import PROGRAM_NAME
+from electos.ballotmaker.data.models import ElectionData
+
+logging.getLogger(__name__)
+
+
+def get_election_header(election: ElectionData) -> dict:
+    """extract the shared data for ballot headers"""
+    name = election.name
+    end_date = election.end_date
+    election_type = election.type
+    return {
+        "Name": name,
+        "EndDate": end_date,
+        "Type": election_type,
+    }
+
+
+def build_ballots(election: ElectionData) -> Path:
+
+    # create the directories needed
+    # TODO: clean up datetime code: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973609852
+    now = datetime.now()
+    date_time = now.strftime("%Y_%m_%dT%H%M%S")
+    home_dir = Path.home()
+    program_dir = Path(home_dir, PROGRAM_NAME)
+    new_ballot_dir = Path(program_dir, date_time)
+    logging.info(f"New ballots will be saved in {new_ballot_dir}")
+    # TODO: Use original EDF file name (no ext) in output dir
+    # See: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973776838
+    Path(new_ballot_dir).mkdir(parents=True, exist_ok=False)
+    # TODO: list actual dir in log output: https://github.com/TrustTheVote-Project/BallotLab/pull/113#discussion_r973610221
+    logging.info("Output directory created.")
+    election_header = get_election_header(election)
+
+    for ballot_data in election.ballot_styles:
+        logging.info(f"Generating ballot for {ballot_data.id}")
+        new_ballot_name = build_ballot(
+            ballot_data, election_header, new_ballot_dir
+        )
+    return new_ballot_dir