<commit_before>import csv
import json
import datetime
import tempfile
import zipfile
import boto3
from django.core.management.base import BaseCommand
from django.db.models import F
from opencivicdata.legislative.models import (
    LegislativeSession,
    Bill,
    BillAbstract,
    BillAction,
    BillTitle,
    BillIdentifier,
    RelatedBill,
    BillSponsorship,
    BillDocument,
    BillVersion,
    BillDocumentLink,
    BillVersionLink,
    BillSource,
    VoteEvent,
    PersonVote,
    VoteCount,
    VoteSource,
)
from ...models import DataExport
from utils.common import abbr_to_jid


def export_csv(filename, data, zf):
    num = len(data)
    if not num:
        return
    headers = data[0].keys()

    with tempfile.NamedTemporaryFile("w") as f:
        print("writing", filename, num, "records")
        of = csv.DictWriter(f, headers)
        of.writeheader()
        of.writerows(data)
        f.flush()
        zf.write(f.name, filename)
    return num


def export_json(filename, data, zf):
    num = len(data)
    if not num:
        return

    with tempfile.NamedTemporaryFile("w") as f:
        print("writing", filename, num, "records")
        json.dump(data, f)
        f.flush()
        zf.write(f.name, filename)
    return num


def export_session_csv(state, session):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    bills = Bill.objects.filter(legislative_session=sobj).values(
        "id",
        "identifier",
        "title",
        "classification",
        "subject",
        session_identifier=F("legislative_session__identifier"),
        jurisdiction=F("legislative_session__jurisdiction__name"),
        organization_classification=F("from_organization__classification"),
    )

    if not bills.count():
        print(f"no bills for {state} {session}")
        return
    filename = f"{state}_{session}_csv.zip"
    zf = zipfile.ZipFile(filename, "w")
    ts = datetime.datetime.utcnow()
    zf.writestr(
        "README",
        f"""Open States Data Export

State: {state}
Session: {session}
Generated At: {ts}
CSV Format Version: 2.0
""",
    )

    export_csv(f"{state}/{session}/{state}_{session}_bills.csv", bills, zf)

    for Model, fname in (
        (BillAbstract, "bill_abstracts"),
        (BillTitle, "bill_titles"),
        (BillIdentifier, "bill_identifiers"),
        (BillAction, "bill_actions"),
        (BillSource, "bill_sources"),
        (RelatedBill, "bill_related_bills"),
        (BillSponsorship, "bill_sponsorships"),
        (BillDocument, "bill_documents"),
        (BillVersion, "bill_versions"),
    ):
        subobjs = Model.objects.filter(bill__legislative_session=sobj).values()
        export_csv(f"{state}/{session}/{state}_{session}_{fname}.csv", subobjs, zf)

    subobjs = BillDocumentLink.objects.filter(
        document__bill__legislative_session=sobj
    ).values()
    export_csv(
        f"{state}/{session}/{state}_{session}_bill_document_links.csv", subobjs, zf
    )
    subobjs = BillVersionLink.objects.filter(
        version__bill__legislative_session=sobj
    ).values()
    export_csv(
        f"{state}/{session}/{state}_{session}_bill_version_links.csv", subobjs, zf
    )

    # TODO: BillActionRelatedEntity

    # Votes
    votes = VoteEvent.objects.filter(legislative_session=sobj).values(
        "id",
        "identifier",
        "motion_text",
        "motion_classification",
        "start_date",
        "result",
        "organization_id",
        "bill_id",
        "bill_action_id",
        jurisdiction=F("legislative_session__jurisdiction__name"),
        session_identifier=F("legislative_session__identifier"),
    )
    export_csv(f"{state}/{session}/{state}_{session}_votes.csv", votes, zf)
    for Model, fname in (
        (PersonVote, "vote_people"),
        (VoteCount, "vote_counts"),
        (VoteSource, "vote_sources"),
    ):
        subobjs = Model.objects.filter(vote_event__legislative_session=sobj).values()
        export_csv(f"{state}/{session}/{state}_{session}_{fname}.csv", subobjs, zf)

    return filename


def export_session_json(state, session):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    bills = list(
        Bill.objects.filter(legislative_session=sobj).values(
            "id",
            "identifier",
            "title",
            "classification",
            "subject",
            session_identifier=F("legislative_session__identifier"),
            jurisdiction=F("legislative_session__jurisdiction__name"),
            organization_classification=F("from_organization__classification"),
            raw_text=F("searchable__raw_text"),
            raw_text_url=F("searchable__version_link__url"),
        )
    )
    filename = f"{state}_{session}_json.zip"
    zf = zipfile.ZipFile(filename, "w")
    ts = datetime.datetime.utcnow()
    zf.writestr(
        "README",
        f"""Open States Data Export

State: {state}
Session: {session}
Generated At: {ts}
JSON Format Version: 0.1
""",
    )

    export_json(f"{state}/{session}/{state}_{session}_bills.json", bills, zf)


def upload_and_publish(state, session, filename, data_type):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    s3 = boto3.client("s3")

    BULK_S3_BUCKET = "data.openstates.org"
    s3_path = f"{data_type}/latest/"
    s3_url = f"https://{BULK_S3_BUCKET}/{s3_path}{filename}"

    s3.upload_file(
        filename, BULK_S3_BUCKET, s3_path + filename, ExtraArgs={"ACL": "public-read"}
    )
    print("uploaded", s3_url)<commit_msg>JSON export format 0.2<commit_after>import csv
import json
import datetime
import tempfile
import zipfile
import boto3
from django.core.management.base import BaseCommand
from django.db.models import F
from opencivicdata.legislative.models import (
    LegislativeSession,
    Bill,
    BillAbstract,
    BillAction,
    BillTitle,
    BillIdentifier,
    RelatedBill,
    BillSponsorship,
    BillDocument,
    BillVersion,
    BillDocumentLink,
    BillVersionLink,
    BillSource,
    VoteEvent,
    PersonVote,
    VoteCount,
    VoteSource,
)
from ...models import DataExport
from utils.common import abbr_to_jid


def export_csv(filename, data, zf):
    num = len(data)
    if not num:
        return
    headers = data[0].keys()

    with tempfile.NamedTemporaryFile("w") as f:
        print("writing", filename, num, "records")
        of = csv.DictWriter(f, headers)
        of.writeheader()
        of.writerows(data)
        f.flush()
        zf.write(f.name, filename)
    return num


def export_json(filename, data, zf):
    num = len(data)
    if not num:
        return

    with tempfile.NamedTemporaryFile("w") as f:
        print("writing", filename, num, "records")
        json.dump(data, f)
        f.flush()
        zf.write(f.name, filename)
    return num


def _docver_to_json(dv):
    return {
        "note": dv.note,
        "date": dv.date,
        "links": list(dv.links.values("url", "media_type")),
    }


def _bill_to_json(b):
    return {
        "id": b.id,
        "legislative_session": b.legislative_session.identifier,
        "jurisdiction_name": b.legislative_session.jurisdiction.name,
        "identifier": b.identifier,
        "title": b.title,
        "chamber": b.from_organization.classification,
        "classification": b.classification,
        "subject": b.subject,
        "abstracts": list(b.abstracts.values("abstract", "note", "date")),
        "other_titles": list(b.other_titles.values("title", "note")),
        "other_identifiers": list(
            b.other_identifiers.values("note", "identifier", "scheme")
        ),
        "actions": list(
            b.actions.values(
                "organization__name", "description", "date", "classification", "order"
            )
        ),
        # TODO: action related entities
        "related_bills": list(b.related_bills.values("related_bill_id")),
        "sponsors": list(b.sponsorships.values("name", "primary", "classification")),
        "documents": [_docver_to_json(d) for d in b.documents.all()],
        "versions": [_docver_to_json(d) for d in b.versions.all()],
        "sources": list(b.sources.values("url")),
        "raw_text": b.searchable.raw_text,
        "raw_text_url": b.searchable.version_link.url
        if b.searchable.version_link
        else None,
    }


def export_session_csv(state, session):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    bills = Bill.objects.filter(legislative_session=sobj).values(
        "id",
        "identifier",
        "title",
        "classification",
        "subject",
        session_identifier=F("legislative_session__identifier"),
        jurisdiction=F("legislative_session__jurisdiction__name"),
        organization_classification=F("from_organization__classification"),
    )

    if not bills.count():
        print(f"no bills for {state} {session}")
        return
    filename = f"{state}_{session}_csv.zip"
    zf = zipfile.ZipFile(filename, "w")
    ts = datetime.datetime.utcnow()
    zf.writestr(
        "README",
        f"""Open States Data Export

State: {state}
Session: {session}
Generated At: {ts}
CSV Format Version: 2.0
""",
    )

    export_csv(f"{state}/{session}/{state}_{session}_bills.csv", bills, zf)

    for Model, fname in (
        (BillAbstract, "bill_abstracts"),
        (BillTitle, "bill_titles"),
        (BillIdentifier, "bill_identifiers"),
        (BillAction, "bill_actions"),
        (BillSource, "bill_sources"),
        (RelatedBill, "bill_related_bills"),
        (BillSponsorship, "bill_sponsorships"),
        (BillDocument, "bill_documents"),
        (BillVersion, "bill_versions"),
    ):
        subobjs = Model.objects.filter(bill__legislative_session=sobj).values()
        export_csv(f"{state}/{session}/{state}_{session}_{fname}.csv", subobjs, zf)

    subobjs = BillDocumentLink.objects.filter(
        document__bill__legislative_session=sobj
    ).values()
    export_csv(
        f"{state}/{session}/{state}_{session}_bill_document_links.csv", subobjs, zf
    )
    subobjs = BillVersionLink.objects.filter(
        version__bill__legislative_session=sobj
    ).values()
    export_csv(
        f"{state}/{session}/{state}_{session}_bill_version_links.csv", subobjs, zf
    )

    # TODO: BillActionRelatedEntity

    # Votes
    votes = VoteEvent.objects.filter(legislative_session=sobj).values(
        "id",
        "identifier",
        "motion_text",
        "motion_classification",
        "start_date",
        "result",
        "organization_id",
        "bill_id",
        "bill_action_id",
        jurisdiction=F("legislative_session__jurisdiction__name"),
        session_identifier=F("legislative_session__identifier"),
    )
    export_csv(f"{state}/{session}/{state}_{session}_votes.csv", votes, zf)
    for Model, fname in (
        (PersonVote, "vote_people"),
        (VoteCount, "vote_counts"),
        (VoteSource, "vote_sources"),
    ):
        subobjs = Model.objects.filter(vote_event__legislative_session=sobj).values()
        export_csv(f"{state}/{session}/{state}_{session}_{fname}.csv", subobjs, zf)

    return filename


def export_session_json(state, session):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    bills = [
        _bill_to_json(b)
        for b in Bill.objects.filter(legislative_session=sobj)
        .select_related(
            "legislative_session",
            "legislative_session__jurisdiction",
            "from_organization",
            "searchable",
        )
        .prefetch_related(
            "abstracts",
            "other_titles",
            "other_identifiers",
            "actions",
            "related_bills",
            "sponsorships",
            "documents",
            "documents__links",
            "versions",
            "versions__links",
            "sources",
        )
    ]
    filename = f"{state}_{session}_json.zip"
    zf = zipfile.ZipFile(filename, "w")
    ts = datetime.datetime.utcnow()
    zf.writestr(
        "README",
        f"""Open States Data Export

State: {state}
Session: {session}
Generated At: {ts}
JSON Format Version: 0.2
""",
    )

    export_json(f"{state}/{session}/{state}_{session}_bills.json", bills, zf)


def upload_and_publish(state, session, filename, data_type):
    sobj = LegislativeSession.objects.get(
        jurisdiction_id=abbr_to_jid(state), identifier=session
    )
    s3 = boto3.client("s3")

    BULK_S3_BUCKET = "data.openstates.org"
    s3_path = f"{data_type}/latest/"
    s3_url = f"https://{BULK_S3_BUCKET}/{s3_path}{filename}"

    s3.upload_file(
        filename, BULK_S3_BUCKET, s3_path + filename, ExtraArgs={"ACL": "public-read"}
    )
    print("uploaded", s3_url)
{'fraction_non_alphanumeric': 0.12706372473383737, 'fraction_numerical': 0.002931646350871779, 'mean_word_length': 2.911587205793603, 'pattern_counts': {'":': 24, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '353618', 'n_tokens_mistral': 4245, 'n_tokens_neox': 4109, 'n_words': 709}