import json
import os

from candidates.models import Ballot
from bulk_adding.models import RawPeople
from django.core.management.base import BaseCommand


class Command(BaseCommand):
    """This command uses the ballots endpoint to loop over each
    ballot and store each sopn pdf (uploaded_file) locally"""

    def add_arguments(self, parser):
        parser.add_argument(
            "--data",
            action="store",
            help="Dictionary of raw people to write as a baseline",
        )

    def handle(self, *args, **options):
        json_data = options["data"] or {}
        if not json_data:
            for ballot in Ballot.objects.exclude(officialdocument__isnull=True):
                raw_people = getattr(ballot, "rawpeople", [])
                try:
                    raw_people = ballot.rawpeople.data
                except RawPeople.DoesNotExist:
                    raw_people = []

                json_data[ballot.ballot_paper_id] = raw_people

        file_path = os.path.join(
            os.getcwd(), "ynr/apps/sopn_parsing/tests/data/sopn_baseline.json"
        )
        with open(file_path, "w") as f:
            f.write(json.dumps(json_data))

@@ -16,16 +16,48 @@
 ort os%0A%0A
+from django.db.models import Q%0A%0A
 from can
@@ -85,16 +85,16 @@
  Ballot%0A
-
 from bul
@@ -219,123 +219,363 @@
  %22%22%22
-This command uses the ballots endpoint to loop over each%0A    ballot and store each sopn pdf (uploaded_file) locally
+%0A    Creates a JSON file to represent ballots that have an Officialdocument.%0A    Only include ballots where:%0A    - The source of the RawPeople is from parsing a PDF%0A    - No RawPeople were created from the OfficialDocument. This is so that we%0A    will know if we make make improvements that mean more RawPeople are parsed%0A    from an OfficialDocument%0A    
 %22%22%22%0A
@@ -852,16 +852,17 @@
 %5D or %7B%7D%0A
+%0A
         
@@ -895,21 +895,12 @@
     
-for ballot in
+qs =
  Bal
@@ -949,16 +949,181 @@
 ll=True)
+.filter(%0A                Q(rawpeople__source_type=RawPeople.SOURCE_PARSED_PDF)%0A                %7C Q(rawpeople__isnull=True)%0A            )%0A            for ballot in qs
 :%0A      

Update which objects are used to write baseline file