<commit_before>"""
Imports ZA provincial and national election candidates using the 2019
IEC spreadsheet format.
"""

import re
import sys
import unicodecsv
import string
import datetime
from optparse import make_option
from pombola.core.models import (
    Organisation,
    OrganisationKind,
    Person,
    Position,
    PositionTitle,
    AlternativePersonName,
)
from django.core.management.base import NoArgsCommand
from django.db.models import Q
from django.db.utils import IntegrityError
from django.utils.text import slugify

from django_date_extensions.fields import ApproximateDate

from haystack.query import SearchQuerySet

# TODO: Check these are the correct dates
final_candidate_list_date = ApproximateDate(year=2019, month=4, day=23)
month_before_final_list_date = ApproximateDate(year=2019, month=3)

party_to_object = {}
list_to_object = {}
position_to_object = {}
YEAR = "2019"
COMMIT = False

csv_files = [
    'National.csv',
    'Provincial-EC.csv',
    'Provincial-FS.csv',
    'Provincial-GP.csv',
    'Provincial-KZN.csv',
    'Provincial-LIM.csv',
    'Provincial-MP.csv',
    'Provincial-NC.csv',
    'Provincial-NW.csv',
    'Provincial-WC.csv',
]

candidates = []

for candidates_csv in csv_files:
    with open("pombola/south_africa/data/elections/2019/" + candidates_csv, "rb") as csvfile:
        csv = unicodecsv.DictReader(csvfile)
        [candidates.append(row) for row in csv]

parties_csv = "pombola/south_africa/data/elections/2019/parties.csv"
party_mapping = {}
with open(parties_csv, "rb") as csvfile:
    csv = unicodecsv.DictReader(csvfile)
    for row in csv:
        party_mapping[row["name_in_candidates_file"]] = row["slug"]

errors = []

def check_or_create_positions():
    """Checks whether positions exist, otherwise create them"""
    num_append = {
        "1": "st",
        "2": "nd",
        "3": "rd",
        "4": "th",
        "5": "th",
        "6": "th",
        "7": "th",
        "8": "th",
        "9": "th",
        "0": "th",
    }
    positions = range(1, 201)

    for p in positions:
        position_name = str(p) + num_append[str(p)[-1]] + " Candidate"
        position_slug = str(p) + num_append[str(p)[-1]] + "_candidate"
        position_object, _ = PositionTitle.objects.get_or_create(
            slug=position_slug, defaults={"name": position_name}
        )
        position_to_object[unicode(p)] = position_object

def get_list(listname, listslug):
    """Returns the organisation object for a list"""
    if listname in list_to_object:
        return list_to_object[listname]

    organisationkind, _ = OrganisationKind.objects.get_or_create(
        slug="election-list", name="Election List"
    )
    listobject, _ = Organisation.objects.get_or_create(
        slug=listslug, name=listname, kind=organisationkind
    )
    list_to_object[listname] = listobject
    return listobject

def get_party(partyname):
    """Returns the organisation object for a party"""
    partyname = re.sub(r"\s+", " ", partyname)
    # if the party has already been worked with, return that
    if partyname in party_to_object:
        return party_to_object[partyname]

    party_slug = party_mapping.get(partyname)
    if not party_slug:
        return False

    party = Organisation.objects.get(slug=party_slug)
    if party:
        party_to_object[partyname] = party
        return party_to_object[partyname]
    else:
        return False

def save_match(
    person,
    party_name,
    list_position,
    import_list_name,
    person_list_firstnames,
    person_list_surname,
):
    """Save imported details about a person who already exists"""
    # get the list to add the person to
    import_list_name = string.replace(import_list_name, ":", "")
    party = get_party(party_name)
    party_parts = string.split(party.name, " (")
    listname = party_parts[0] + " " + import_list_name + " Election List " + YEAR
    listslug = (
        party.slug
        + "-"
        + string.replace(import_list_name.lower(), " ", "-")
        + "-election-list-"
        + YEAR
    )
    list_object = get_list(listname, listslug)

    # get the position title
    positiontitle = position_to_object[list_position]

    # check if we have distinct family/given names on record - if we do
    # we assume these are absolutely correct and so if they do not match
    # the import values the match is considered incorrect
    if person.given_name == "" and person.family_name == "":
        person.given_name = person_list_firstnames.title()
        person.family_name = person_list_surname.title()
        if COMMIT:
            print "- updating given_name/family_name"
            person.save()
        else:
            print "* would update blank given_name/family_name"
    elif string.replace(person_list_firstnames.lower(), "-", " ") != string.replace(
        person.given_name.lower(), "-", " "
    ) or string.replace(person_list_surname.lower(), "-", " ") != string.replace(
        person.family_name.lower(), "-", " "
    ):
        print "- but names do not match existing given_name/family_name, ignoring"
        return False

    # get the person's current party to compare whether these have changed
    # we assume that a person on an election list should not be a member
    # more than one party - so we end memberships that do not match the list
    foundcorrectparty = False
    for p in person.parties():
        if p == party:
            foundcorrectparty = True
        else:
            # end the membership
            now = datetime.date.today()
            now_approx = repr(
                ApproximateDate(year=now.year, month=now.month, day=now.day)
            )
            position = Position.objects.get(
                Q(person=person),
                Q(organisation=p),
                Q(title__slug="member"),
                Q(start_date__lte=now_approx),
                (Q(sorting_end_date_high__gte=now_approx) | Q(end_date="")),
            )
            position.end_date = month_before_final_list_date
            if COMMIT:
                print "- ending %s party position" % position.organisation
                position.save()
            else:
                print "* would end %s party position" % position.organisation

    if COMMIT and not foundcorrectparty:
        print "- creating %s party position" % party
        # get the member title
        member = PositionTitle.objects.get(name="Member")
        # add the correct membership
        position = Position.objects.get_or_create(
            person=person,
            organisation=party,
            start_date=month_before_final_list_date,
            title=member,
            category="political",
        )

    # check whether the list name matches the recorded name
    if (
        person_list_firstnames + " " + person_list_surname
    ).lower() != person.legal_name.lower():
        # set the new name as the person name and record the old as an
        # alternative name
        new_name = (person_list_firstnames + " " + person_list_surname).title()
        if COMMIT:
            print "- updating person name from %s to %s" % (person.legal_name, new_name)
            AlternativePersonName.objects.get_or_create(
                person=person, alternative_name=person.name
            )
            person.legal_name = new_name
            person.save()
        else:
            print "* would update person legal_name from %s to %s" % (
                person.legal_name,
                new_name,
            )

    # add the actual membership of the list organisation. Note no end date
    # is set at this point. This should probably be the date of the election
    # but it is unclear whether setting this immediately would cause issues
    # if viewed on the election date or immediately afterwards
    if COMMIT:
        print "- creating candidate position"
        position = Position.objects.get_or_create(
            person=person,
            organisation=list_object,
            start_date=final_candidate_list_date,
            title=positiontitle,
            category="political",
        )

    return True

def add_new_person(
    party_name,
    list_position,
    import_list_name,
    person_list_firstnames,
    person_list_surname,
):
    """Create a new person with appropriate positions"""

    if not COMMIT:
        print "would create new entry"
        return

    print "creating"
    # get the list to add the person to
    import_list_name = string.replace(import_list_name, ":", "")
    party = get_party(party_name)
    party_parts = string.split(party.name, " (")
    listname = party_parts[0] + " " + import_list_name + " Election List " + YEAR
    listslug = (
        party.slug
        + "-"
        + string.replace(import_list_name.lower(), " ", "-")
        + "-election-list-"
        + YEAR
    )
    list_object = get_list(listname, listslug)

    # get the position title
    positiontitle = position_to_object[list_position]

    slug = slugify(person_list_firstnames + " " + person_list_surname)

    try:
        # create the person
        person, _ = Person.objects.get_or_create(
            legal_name=(person_list_firstnames + " " + person_list_surname).title(),
            given_name=person_list_firstnames.title(),
            family_name=person_list_surname.title(),
            slug=slug,
        )
    except IntegrityError as e:
        errors.append(e)
        person = Person.objects.get(slug=slug)

    # add to the party
    member = PositionTitle.objects.get(name="Member")
    position, _ = Position.objects.get_or_create(
        person=person, organisation=party, title=member, category="political"
    )

    # add the actual membership of the list organisation. Note no end date
    # is set at this point. This should probably be the date of the election
    # but it is unclear whether setting this immediately would cause issues
    # if viewed on the election date or immediately afterwards
    position, _ = Position.objects.get_or_create(
        person=person,
        organisation=list_object,
        start_date=final_candidate_list_date,
        title=positiontitle,
        category="political",
    )

def process_search(firstnames, surname, party, list_position, list_name, url):
    """Process a search based on the search string passed via url"""
    search = SearchQuerySet().models(Person).filter(text=url)
    if len(search) > 1:
        name = (firstnames + " " + surname).title()
        search = [s for s in search if s.object and s.object.legal_name == name]
    if len(search) == 1 and search[0].object:
        print "match", search[0].object.name
        return save_match(
            search[0].object, party, list_position, list_name, firstnames, surname
        )
    else:
        return False

# Search 1
def search_full_name(firstnames, surname, party, list_position, list_name):
    searchurl = '"' + (firstnames + " " + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 2
def search_reordered(firstnames, surname, party, list_position, list_name):
    searchurl = '"' + (firstnames + " " + surname) + '"~4'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 3
def search_first_names(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    if len(names) == 1:
        # Already covered by search 1
        return False
    for name in names:
        searchurl = '"' + (name + " " + surname) + '"'
        if process_search(
            firstnames, surname, party, list_position, list_name, searchurl
        ):
            return True
    return False

# Search 4
def search_initials(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    initials = ""
    for name in names:
        if name[:1] == "A":
            continue
        initials += name[:1] + ". "
    if not initials:
        return False
    searchurl = '"' + (initials + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 5
def search_initials_alt(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    initials = ""
    for name in names:
        if name[:1] == "A":
            continue
        initials += name[:1] + "."
    if not initials:
        return False
    searchurl = '"' + (initials + " " + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 6
def search_misspellings(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    first = ""
    for name in names:
        first += name + "~ AND "
    searchurl = first + surname
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

def search(firstnames, surname, party, list_position, list_name):
    """Attempt varius approaches to matching names"""
    print "Looking at %s %s:" % (firstnames, surname),

    existing = Person.objects.filter(
        given_name__iexact=firstnames, family_name__iexact=surname
    )
    if len(existing) == 1:
        print "match existing", existing[0].name
        if save_match(
            existing[0], party, list_position, list_name, firstnames, surname
        ):
            return True

    if search_full_name(firstnames, surname, party, list_position, list_name):
        return True
    if search_reordered(firstnames, surname, party, list_position, list_name):
        return True
    if search_first_names(firstnames, surname, party, list_position, list_name):
        return True
    if search_initials(firstnames, surname, party, list_position, list_name):
        return True
    if search_initials_alt(firstnames, surname, party, list_position, list_name):
        return True
    if search_misspellings(firstnames, surname, party, list_position, list_name):
        return True
    return False

class Command(NoArgsCommand):
    """Import South African national and provincial election candidates"""

    help = (
        "Import csv file of South African national and provincial election candidates"
    )

    option_list = NoArgsCommand.option_list + (
        make_option(
            "--commit",
            action="store_true",
            help="Actually commit person changes to the database (new positions/orgs always created)",
        ),
    )

    def handle_noargs(self, **options):
        global COMMIT
        COMMIT = options["commit"]

        # check all the parties exist
        missingparties = set()
        for row in candidates:
            party_name = row["Party name"]
            if not get_party(party_name):
                missingparties.add(party_name)
        if missingparties:
            for party in sorted(missingparties):
                print "Missing party:", party
            sys.exit(1)

        # check whether the positions exist, otherwise create them
        check_or_create_positions()

        for row in candidates:
            party_name = row["Party name"].strip()
            list_type = row["List type"].strip()
            order_number = row["Order number"].strip()
            full_names = row["Full names"].strip()
            surname = row["Surname"].strip()

            if not search(full_names, surname, party_name, order_number, list_type):
                add_new_person(party_name, order_number, list_type, full_names, surname)

        if errors:
            print("ERRORS:")
            for error in errors:
                print(error)
<commit_msg>[ZA] Remove TODO<commit_after>"""
Imports ZA provincial and national election candidates using the 2019
IEC spreadsheet format.
"""

import re
import sys
import unicodecsv
import string
import datetime
from optparse import make_option
from pombola.core.models import (
    Organisation,
    OrganisationKind,
    Person,
    Position,
    PositionTitle,
    AlternativePersonName,
)
from django.core.management.base import NoArgsCommand
from django.db.models import Q
from django.db.utils import IntegrityError
from django.utils.text import slugify

from django_date_extensions.fields import ApproximateDate

from haystack.query import SearchQuerySet

final_candidate_list_date = ApproximateDate(year=2019, month=4, day=23)
month_before_final_list_date = ApproximateDate(year=2019, month=3)

party_to_object = {}
list_to_object = {}
position_to_object = {}
YEAR = "2019"
COMMIT = False

csv_files = [
    'National.csv',
    'Provincial-EC.csv',
    'Provincial-FS.csv',
    'Provincial-GP.csv',
    'Provincial-KZN.csv',
    'Provincial-LIM.csv',
    'Provincial-MP.csv',
    'Provincial-NC.csv',
    'Provincial-NW.csv',
    'Provincial-WC.csv',
]

candidates = []

for candidates_csv in csv_files:
    with open("pombola/south_africa/data/elections/2019/" + candidates_csv, "rb") as csvfile:
        csv = unicodecsv.DictReader(csvfile)
        [candidates.append(row) for row in csv]

parties_csv = "pombola/south_africa/data/elections/2019/parties.csv"
party_mapping = {}
with open(parties_csv, "rb") as csvfile:
    csv = unicodecsv.DictReader(csvfile)
    for row in csv:
        party_mapping[row["name_in_candidates_file"]] = row["slug"]

errors = []

def check_or_create_positions():
    """Checks whether positions exist, otherwise create them"""
    num_append = {
        "1": "st",
        "2": "nd",
        "3": "rd",
        "4": "th",
        "5": "th",
        "6": "th",
        "7": "th",
        "8": "th",
        "9": "th",
        "0": "th",
    }
    positions = range(1, 201)

    for p in positions:
        position_name = str(p) + num_append[str(p)[-1]] + " Candidate"
        position_slug = str(p) + num_append[str(p)[-1]] + "_candidate"
        position_object, _ = PositionTitle.objects.get_or_create(
            slug=position_slug, defaults={"name": position_name}
        )
        position_to_object[unicode(p)] = position_object

def get_list(listname, listslug):
    """Returns the organisation object for a list"""
    if listname in list_to_object:
        return list_to_object[listname]

    organisationkind, _ = OrganisationKind.objects.get_or_create(
        slug="election-list", name="Election List"
    )
    listobject, _ = Organisation.objects.get_or_create(
        slug=listslug, name=listname, kind=organisationkind
    )
    list_to_object[listname] = listobject
    return listobject

def get_party(partyname):
    """Returns the organisation object for a party"""
    partyname = re.sub(r"\s+", " ", partyname)
    # if the party has already been worked with, return that
    if partyname in party_to_object:
        return party_to_object[partyname]

    party_slug = party_mapping.get(partyname)
    if not party_slug:
        return False

    party = Organisation.objects.get(slug=party_slug)
    if party:
        party_to_object[partyname] = party
        return party_to_object[partyname]
    else:
        return False

def save_match(
    person,
    party_name,
    list_position,
    import_list_name,
    person_list_firstnames,
    person_list_surname,
):
    """Save imported details about a person who already exists"""
    # get the list to add the person to
    import_list_name = string.replace(import_list_name, ":", "")
    party = get_party(party_name)
    party_parts = string.split(party.name, " (")
    listname = party_parts[0] + " " + import_list_name + " Election List " + YEAR
    listslug = (
        party.slug
        + "-"
        + string.replace(import_list_name.lower(), " ", "-")
        + "-election-list-"
        + YEAR
    )
    list_object = get_list(listname, listslug)

    # get the position title
    positiontitle = position_to_object[list_position]

    # check if we have distinct family/given names on record - if we do
    # we assume these are absolutely correct and so if they do not match
    # the import values the match is considered incorrect
    if person.given_name == "" and person.family_name == "":
        person.given_name = person_list_firstnames.title()
        person.family_name = person_list_surname.title()
        if COMMIT:
            print "- updating given_name/family_name"
            person.save()
        else:
            print "* would update blank given_name/family_name"
    elif string.replace(person_list_firstnames.lower(), "-", " ") != string.replace(
        person.given_name.lower(), "-", " "
    ) or string.replace(person_list_surname.lower(), "-", " ") != string.replace(
        person.family_name.lower(), "-", " "
    ):
        print "- but names do not match existing given_name/family_name, ignoring"
        return False

    # get the person's current party to compare whether these have changed
    # we assume that a person on an election list should not be a member
    # more than one party - so we end memberships that do not match the list
    foundcorrectparty = False
    for p in person.parties():
        if p == party:
            foundcorrectparty = True
        else:
            # end the membership
            now = datetime.date.today()
            now_approx = repr(
                ApproximateDate(year=now.year, month=now.month, day=now.day)
            )
            position = Position.objects.get(
                Q(person=person),
                Q(organisation=p),
                Q(title__slug="member"),
                Q(start_date__lte=now_approx),
                (Q(sorting_end_date_high__gte=now_approx) | Q(end_date="")),
            )
            position.end_date = month_before_final_list_date
            if COMMIT:
                print "- ending %s party position" % position.organisation
                position.save()
            else:
                print "* would end %s party position" % position.organisation

    if COMMIT and not foundcorrectparty:
        print "- creating %s party position" % party
        # get the member title
        member = PositionTitle.objects.get(name="Member")
        # add the correct membership
        position = Position.objects.get_or_create(
            person=person,
            organisation=party,
            start_date=month_before_final_list_date,
            title=member,
            category="political",
        )

    # check whether the list name matches the recorded name
    if (
        person_list_firstnames + " " + person_list_surname
    ).lower() != person.legal_name.lower():
        # set the new name as the person name and record the old as an
        # alternative name
        new_name = (person_list_firstnames + " " + person_list_surname).title()
        if COMMIT:
            print "- updating person name from %s to %s" % (person.legal_name, new_name)
            AlternativePersonName.objects.get_or_create(
                person=person, alternative_name=person.name
            )
            person.legal_name = new_name
            person.save()
        else:
            print "* would update person legal_name from %s to %s" % (
                person.legal_name,
                new_name,
            )

    # add the actual membership of the list organisation. Note no end date
    # is set at this point. This should probably be the date of the election
    # but it is unclear whether setting this immediately would cause issues
    # if viewed on the election date or immediately afterwards
    if COMMIT:
        print "- creating candidate position"
        position = Position.objects.get_or_create(
            person=person,
            organisation=list_object,
            start_date=final_candidate_list_date,
            title=positiontitle,
            category="political",
        )

    return True

def add_new_person(
    party_name,
    list_position,
    import_list_name,
    person_list_firstnames,
    person_list_surname,
):
    """Create a new person with appropriate positions"""

    if not COMMIT:
        print "would create new entry"
        return

    print "creating"
    # get the list to add the person to
    import_list_name = string.replace(import_list_name, ":", "")
    party = get_party(party_name)
    party_parts = string.split(party.name, " (")
    listname = party_parts[0] + " " + import_list_name + " Election List " + YEAR
    listslug = (
        party.slug
        + "-"
        + string.replace(import_list_name.lower(), " ", "-")
        + "-election-list-"
        + YEAR
    )
    list_object = get_list(listname, listslug)

    # get the position title
    positiontitle = position_to_object[list_position]

    slug = slugify(person_list_firstnames + " " + person_list_surname)

    try:
        # create the person
        person, _ = Person.objects.get_or_create(
            legal_name=(person_list_firstnames + " " + person_list_surname).title(),
            given_name=person_list_firstnames.title(),
            family_name=person_list_surname.title(),
            slug=slug,
        )
    except IntegrityError as e:
        errors.append(e)
        person = Person.objects.get(slug=slug)

    # add to the party
    member = PositionTitle.objects.get(name="Member")
    position, _ = Position.objects.get_or_create(
        person=person, organisation=party, title=member, category="political"
    )

    # add the actual membership of the list organisation. Note no end date
    # is set at this point. This should probably be the date of the election
    # but it is unclear whether setting this immediately would cause issues
    # if viewed on the election date or immediately afterwards
    position, _ = Position.objects.get_or_create(
        person=person,
        organisation=list_object,
        start_date=final_candidate_list_date,
        title=positiontitle,
        category="political",
    )

def process_search(firstnames, surname, party, list_position, list_name, url):
    """Process a search based on the search string passed via url"""
    search = SearchQuerySet().models(Person).filter(text=url)
    if len(search) > 1:
        name = (firstnames + " " + surname).title()
        search = [s for s in search if s.object and s.object.legal_name == name]
    if len(search) == 1 and search[0].object:
        print "match", search[0].object.name
        return save_match(
            search[0].object, party, list_position, list_name, firstnames, surname
        )
    else:
        return False

# Search 1
def search_full_name(firstnames, surname, party, list_position, list_name):
    searchurl = '"' + (firstnames + " " + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 2
def search_reordered(firstnames, surname, party, list_position, list_name):
    searchurl = '"' + (firstnames + " " + surname) + '"~4'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 3
def search_first_names(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    if len(names) == 1:
        # Already covered by search 1
        return False
    for name in names:
        searchurl = '"' + (name + " " + surname) + '"'
        if process_search(
            firstnames, surname, party, list_position, list_name, searchurl
        ):
            return True
    return False

# Search 4
def search_initials(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    initials = ""
    for name in names:
        if name[:1] == "A":
            continue
        initials += name[:1] + ". "
    if not initials:
        return False
    searchurl = '"' + (initials + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 5
def search_initials_alt(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    initials = ""
    for name in names:
        if name[:1] == "A":
            continue
        initials += name[:1] + "."
    if not initials:
        return False
    searchurl = '"' + (initials + " " + surname) + '"'
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

# Search 6
def search_misspellings(firstnames, surname, party, list_position, list_name):
    names = firstnames.split(" ")
    first = ""
    for name in names:
        first += name + "~ AND "
    searchurl = first + surname
    return process_search(
        firstnames, surname, party, list_position, list_name, searchurl
    )

def search(firstnames, surname, party, list_position, list_name):
    """Attempt varius approaches to matching names"""
    print "Looking at %s %s:" % (firstnames, surname),

    existing = Person.objects.filter(
        given_name__iexact=firstnames, family_name__iexact=surname
    )
    if len(existing) == 1:
        print "match existing", existing[0].name
        if save_match(
            existing[0], party, list_position, list_name, firstnames, surname
        ):
            return True

    if search_full_name(firstnames, surname, party, list_position, list_name):
        return True
    if search_reordered(firstnames, surname, party, list_position, list_name):
        return True
    if search_first_names(firstnames, surname, party, list_position, list_name):
        return True
    if search_initials(firstnames, surname, party, list_position, list_name):
        return True
    if search_initials_alt(firstnames, surname, party, list_position, list_name):
        return True
    if search_misspellings(firstnames, surname, party, list_position, list_name):
        return True
    return False

class Command(NoArgsCommand):
    """Import South African national and provincial election candidates"""

    help = (
        "Import csv file of South African national and provincial election candidates"
    )

    option_list = NoArgsCommand.option_list + (
        make_option(
            "--commit",
            action="store_true",
            help="Actually commit person changes to the database (new positions/orgs always created)",
        ),
    )

    def handle_noargs(self, **options):
        global COMMIT
        COMMIT = options["commit"]

        # check all the parties exist
        missingparties = set()
        for row in candidates:
            party_name = row["Party name"]
            if not get_party(party_name):
                missingparties.add(party_name)
        if missingparties:
            for party in sorted(missingparties):
                print "Missing party:", party
            sys.exit(1)

        # check whether the positions exist, otherwise create them
        check_or_create_positions()

        for row in candidates:
            party_name = row["Party name"].strip()
            list_type = row["List type"].strip()
            order_number = row["Order number"].strip()
            full_names = row["Full names"].strip()
            surname = row["Surname"].strip()

            if not search(full_names, surname, party_name, order_number, list_type):
                add_new_person(party_name, order_number, list_type, full_names, surname)

        if errors:
            print("ERRORS:")
            for error in errors:
                print(error)
<|endoftext|>


<commit_before>"""
41.25%
87.54%
"""

class Solution(object):
    # def findDuplicates(self, nums):
    #     """
    #     :type nums: List[int]
    #     :rtype: List[int]
    #     """

    #     mp = {}
    #     result = []
    #     for n in nums:
    #         try:
    #             mp[n] += 1
    #         except:
    #             mp[n] = 1
    #         else:
    #             result.append(n)
    #     return result

    def findDuplicates(self, nums):
        """
        :type nums: List[int]
        :rtype: List[int]
        """
        result = []
        for n in nums:
            if nums[abs(n)-1] < 0:
                result.append(abs(n))
            else:
                nums[abs(n)-1] *= -1
        return result<commit_msg>no message<commit_after>"""
41.25%
87.54%
"""

class Solution(object):
    # def findDuplicates(self, nums):
    #     """
    #     :type nums: List[int]
    #     :rtype: List[int]
    #     """

    #     mp = {}
    #     result = []
    #     for n in nums:
    #         try:
    #             mp[n] += 1
    #         except:
    #             mp[n] = 1
    #         else:
    #             result.append(n)
    #     return result

    def findDuplicates(self, nums):
        """
        O(1) space not including the input and output variables

The idea is we do a linear pass using the input array itself as a hash to store which numbers have been seen before. We do this by making elements at certain indexes negative. See the full explanation here

http://www.geeksforgeeks.org/find-duplicates-in-on-time-and-constant-extra-space/

        :type nums: List[int]
        :rtype: List[int]
        """
        result = []
        for n in nums:
            if nums[abs(n)-1] < 0:
                result.append(abs(n))
            else:
                nums[abs(n)-1] *= -1
        return result<|endoftext|>


<commit_before>from gensim.models import word2vec

class W2CFeatureGenerator:

    def __init__(self):
        # Set values for various parameters
        self.num_features = 5 # Word vector dimensionality
        self.min_word_count = 1  # Minimum word count
        self.num_workers = 4  # Number of threads to run in parallel
        self.context = 10  # Context window size
        self.down_sampling = 1e-3  # Downsample setting for frequent words

    def generate_features(self,sentences):
        # Initialize and train the model (this will take some time)
        print("Training model...")
        model = word2vec.Word2Vec(sentences, workers=self.num_workers, \
                                  size=self.num_features, min_count=self.min_word_count, \
                                  window=self.context, sample=self.down_sampling)

        # If you don't plan to train the model any further, calling
        # init_sims will make the model much more memory-efficient.
        model.init_sims(replace=True)

        # It can be helpful to create a meaningful model name and
        # save the model for later use. You can load it later using Word2Vec.load()
        model_name = "classifier"
        model.save('model/'+model_name)
        print("Keys after training - " + str(list(model.wv.vocab.keys())))
        print("Shape - " + str(model.wv.syn0.shape))
        print(model.wv.similarity("monei","payment"))
        return model<commit_msg>Delete w2c_generator.py<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Licensed under a MIT style license - see LICENSE.rst

"""Generate paths to BOSS data files.

The path module provides convenience methods for building the paths of frequently used
data files.  Most scripts will create a single :class:`Finder` object using the default
constructor for this purpose::

    import bossdata.path
    finder = bossdata.path.Finder()

This finder object is configured by the `$BOSS_SAS_PATH` and `$BOSS_REDUX_VERSION`
environment variables and no other modules uses these variables, except through a
a :class:`Finder` object.  :class:`Finder` objects never interact with any local or
remote filesystems: use the :mod:`bossdata.remote` module to download data files
and access them locally. See :doc:`/usage` for recommendations
on using the :mod:`bossdata.path` and :mod:`bossdata.remote` modules together.
"""

from __future__ import division, print_function

import os
import os.path

class Finder(object):
    """Initialize a path finder object.

    When the constructor is called with no arguments, it will raise a ValueError if either
    BOSS_SAS_PATH or BOSS_REDUX_VERSION is not set.

    Args:
        sas_root(str): Location of the SAS root path to use, e.g., /sas/dr12. Will use the
            value of the BOSS_SAS_PATH environment variable if this is not set.
        redux_version(str): String tag specifying the BOSS spectro reduction version to use,
            e.g., v5_7_0. Will use the value of the BOSS_REDUX_VERSION environment variable
            if this is not set.

    Raises:
        ValueError: No SAS root or redux version specified on the command line or via
            environment variables.
    """
    def __init__(self, sas_root=None, redux_version=None):
        if sas_root is None:
            sas_root = os.getenv('BOSS_SAS_PATH')
        if sas_root is None:
            raise ValueError('No SAS root specified: try setting $BOSS_SAS_PATH.')

        if redux_version is None:
            redux_version = os.getenv('BOSS_REDUX_VERSION', None)
        if redux_version is None:
            raise ValueError('No redux version specifed: try setting $BOSS_REDUX_VERSION.')

        self.redux_version = redux_version
        self.redux_base = os.path.join(sas_root, 'spectro', 'redux', redux_version)

    def get_plate_path(self, plate):
        """Get the path to the specified plate.

        The returned path contains files that include all targets on the plate. Use the
        :meth:`get_spec_path` method for the path of a single spectrum file.

        This method only performs minimal checks that the requested plate number is valid.

        Args:
            plate(int): Plate number, which must be positive.

        Returns:
            str: Full path to files for the specified plate.

        Raises:
            ValueError: Invalid plate number must be > 0.
        """
        if plate < 0:
            raise ValueError('Invalid plate number ({}) must be > 0.'.format(plate))

        return os.path.join(self.redux_base, str(plate))

    def get_plate_plan_path(self, plate, mjd, combined=True):
        """Get the path to the specified plate plan file.

        A combined plan may span several nearby MJDs, in which case the last MJD is
        the one used to identify the plan.

        Args:
            plate(int): Plate number, which must be positive.
            mjd(int): Modified Julian date of the observation, which must be > 3500.
            combined(bool): Specifies the combined plan, which spans all MJDs
                associated with a coadd, but does not include calibration frames
                (arcs,flats) for a specific MJD.

        Returns:
            str: Full path to the requested plan file.

        Raises:
            ValueError: Invalid plate or mjd inputs.
        """
        if mjd <= 3500:
            raise ValueError('Invalid mjd ({}) must be >= 3500.'.format(mjd))

        if combined:
            filename = 'spPlancomb-{plate:4d}-{mjd:5d}.par'.format(plate=plate, mjd=mjd)
        else:
            filename = 'spPlan2d-{plate:4d}-{mjd:5d}.par'.format(plate=plate, mjd=mjd)
        return os.path.join(self.get_plate_path(plate), filename)

    def get_sp_all_path(self, lite=True):
        """Get the location of the metadata summary file.

        The data model of the full (non-lite) file is at
        http://dr12.sdss3.org/datamodel/files/BOSS_SPECTRO_REDUX/RUN2D/spAll.html
        As of DR12, the full file size is about 10Gb and the lite file is about 115Mb.

        Args:
            lite(bool): Specifies the "lite" version which contains all rows but only the
                most commonly used subset of columns. The lite version is a compressed (.gz)
                text data file, while the full version is a FITS file.
        """
        if lite:
            name = 'spAll-{}.dat.gz'.format(self.redux_version)
        else:
            name = 'spAll-{}.fits'.format(self.redux_version)
        return os.path.join(self.redux_base, name)

    def get_spec_path(self, plate, mjd, fiber, lite=True):
        """Get the location of the spectrum file for the specified observation.

        The DR12 data model for the returned files is at
        http://dr12.sdss3.org/datamodel/files/BOSS_SPECTRO_REDUX/RUN2D/spectra/PLATE4/spec.html
        but only HDUs 0-3 are included in the (default) lite format.
        Each lite (full) file is approximately 0.2Mb (1.7Mb) in size.

        Use the :meth:`get_plate_path` method for the path to files that include all targets
        on a plate.

        This method only performs minimal checks that the requested plate-mjd-fiber are valid.

        Args:
            plate(int): Plate number, which must be positive.
            mjd(int): Modified Julian date of the observation, which must be > 3500.
            fiber(int): Fiber number of the target on this plate, which must be in the
                range 1-1000.
            lite(bool): Specifies the "lite" version which contains only HDUs 0-3, so no
                per-exposure data is included.

        Returns:
            str: Full path to the spectrum file for the specified observation.

        Raises:
            ValueError: Invalid plate, mjd or fiber inputs.
        """
        if plate < 0:
            raise ValueError('Invalid plate number ({}) must be > 0.'.format(plate))
        if mjd <= 3500:
            raise ValueError('Invalid mjd ({}) must be >= 3500.'.format(mjd))
        if fiber < 1 or fiber > 1000:
            raise ValueError('Invalid fiber ({}) must be 1-1000.'.format(fiber))

        path = 'spectra'
        if lite:
            path = os.path.join(path, 'lite')
        name = 'spec-{plate:4d}-{mjd:5d}-{fiber:04d}.fits'.format(
            plate=plate, mjd=mjd, fiber=fiber)
        return os.path.join(self.redux_base, path, str(plate), name)
<commit_msg>add get_quasar_catalog_path() function to Finder class<commit_after># -*- coding: utf-8 -*-
# Licensed under a MIT style license - see LICENSE.rst

"""Generate paths to BOSS data files.

The path module provides convenience methods for building the paths of frequently used
data files.  Most scripts will create a single :class:`Finder` object using the default
constructor for this purpose::

    import bossdata.path
    finder = bossdata.path.Finder()

This finder object is configured by the `$BOSS_SAS_PATH` and `$BOSS_REDUX_VERSION`
environment variables and no other modules uses these variables, except through a
a :class:`Finder` object.  :class:`Finder` objects never interact with any local or
remote filesystems: use the :mod:`bossdata.remote` module to download data files
and access them locally. See :doc:`/usage` for recommendations
on using the :mod:`bossdata.path` and :mod:`bossdata.remote` modules together.
"""

from __future__ import division, print_function

import os
import os.path

class Finder(object):
    """Initialize a path finder object.

    When the constructor is called with no arguments, it will raise a ValueError if either
    BOSS_SAS_PATH or BOSS_REDUX_VERSION is not set.

    Args:
        sas_root(str): Location of the SAS root path to use, e.g., /sas/dr12. Will use the
            value of the BOSS_SAS_PATH environment variable if this is not set.
        redux_version(str): String tag specifying the BOSS spectro reduction version to use,
            e.g., v5_7_0. Will use the value of the BOSS_REDUX_VERSION environment variable
            if this is not set.

    Raises:
        ValueError: No SAS root or redux version specified on the command line or via
            environment variables.
    """
    def __init__(self, sas_root=None, redux_version=None):
        if sas_root is None:
            sas_root = os.getenv('BOSS_SAS_PATH')
        if sas_root is None:
            raise ValueError('No SAS root specified: try setting $BOSS_SAS_PATH.')

        if redux_version is None:
            redux_version = os.getenv('BOSS_REDUX_VERSION', None)
        if redux_version is None:
            raise ValueError('No redux version specifed: try setting $BOSS_REDUX_VERSION.')

        self.sas_root = sas_root
        self.redux_version = redux_version
        self.redux_base = os.path.join(sas_root, 'spectro', 'redux', redux_version)

    def get_plate_path(self, plate):
        """Get the path to the specified plate.

        The returned path contains files that include all targets on the plate. Use the
        :meth:`get_spec_path` method for the path of a single spectrum file.

        This method only performs minimal checks that the requested plate number is valid.

        Args:
            plate(int): Plate number, which must be positive.

        Returns:
            str: Full path to files for the specified plate.

        Raises:
            ValueError: Invalid plate number must be > 0.
        """
        if plate < 0:
            raise ValueError('Invalid plate number ({}) must be > 0.'.format(plate))

        return os.path.join(self.redux_base, str(plate))

    def get_plate_plan_path(self, plate, mjd, combined=True):
        """Get the path to the specified plate plan file.

        A combined plan may span several nearby MJDs, in which case the last MJD is
        the one used to identify the plan.

        Args:
            plate(int): Plate number, which must be positive.
            mjd(int): Modified Julian date of the observation, which must be > 3500.
            combined(bool): Specifies the combined plan, which spans all MJDs
                associated with a coadd, but does not include calibration frames
                (arcs,flats) for a specific MJD.

        Returns:
            str: Full path to the requested plan file.

        Raises:
            ValueError: Invalid plate or mjd inputs.
        """
        if mjd <= 3500:
            raise ValueError('Invalid mjd ({}) must be >= 3500.'.format(mjd))

        if combined:
            filename = 'spPlancomb-{plate:4d}-{mjd:5d}.par'.format(plate=plate, mjd=mjd)
        else:
            filename = 'spPlan2d-{plate:4d}-{mjd:5d}.par'.format(plate=plate, mjd=mjd)
        return os.path.join(self.get_plate_path(plate), filename)

    def get_sp_all_path(self, lite=True):
        """Get the location of the metadata summary file.

        The data model of the full (non-lite) file is at
        http://dr12.sdss3.org/datamodel/files/BOSS_SPECTRO_REDUX/RUN2D/spAll.html
        As of DR12, the full file size is about 10Gb and the lite file is about 115Mb.

        Args:
            lite(bool): Specifies the "lite" version which contains all rows but only the
                most commonly used subset of columns. The lite version is a compressed (.gz)
                text data file, while the full version is a FITS file.
        """
        if lite:
            name = 'spAll-{}.dat.gz'.format(self.redux_version)
        else:
            name = 'spAll-{}.fits'.format(self.redux_version)
        return os.path.join(self.redux_base, name)

    def get_quasar_catalog_path(self):
        """Get the location of the quasar catalog file.

        The data model of the quasar catalog is at
        http://data.sdss3.org/datamodel/files/BOSS_QSO/DR12Q/DR12Q.html
        As of DR12, the file size is about 513Mb.

        Additional information is available at
        http://www.sdss.org/dr12/algorithms/boss-dr12-quasar-catalog/
        """
        # How do we want to specify this?
        #    function arg, env var, guess from sas_path?
        catalog_name = 'DR12Q'
        filename = '{}.fits'.format(catalog_name)
        return os.path.join(self.sas_root, 'qso', catalog_name, filename)

    def get_spec_path(self, plate, mjd, fiber, lite=True):
        """Get the location of the spectrum file for the specified observation.

        The DR12 data model for the returned files is at
        http://dr12.sdss3.org/datamodel/files/BOSS_SPECTRO_REDUX/RUN2D/spectra/PLATE4/spec.html
        but only HDUs 0-3 are included in the (default) lite format.
        Each lite (full) file is approximately 0.2Mb (1.7Mb) in size.

        Use the :meth:`get_plate_path` method for the path to files that include all targets
        on a plate.

        This method only performs minimal checks that the requested plate-mjd-fiber are valid.

        Args:
            plate(int): Plate number, which must be positive.
            mjd(int): Modified Julian date of the observation, which must be > 3500.
            fiber(int): Fiber number of the target on this plate, which must be in the
                range 1-1000.
            lite(bool): Specifies the "lite" version which contains only HDUs 0-3, so no
                per-exposure data is included.

        Returns:
            str: Full path to the spectrum file for the specified observation.

        Raises:
            ValueError: Invalid plate, mjd or fiber inputs.
        """
        if plate < 0:
            raise ValueError('Invalid plate number ({}) must be > 0.'.format(plate))
        if mjd <= 3500:
            raise ValueError('Invalid mjd ({}) must be >= 3500.'.format(mjd))
        if fiber < 1 or fiber > 1000:
            raise ValueError('Invalid fiber ({}) must be 1-1000.'.format(fiber))

        path = 'spectra'
        if lite:
            path = os.path.join(path, 'lite')
        name = 'spec-{plate:4d}-{mjd:5d}-{fiber:04d}.fits'.format(
            plate=plate, mjd=mjd, fiber=fiber)
        return os.path.join(self.redux_base, path, str(plate), name)
<|endoftext|>


<commit_before>"""
 1. Install Scruffy
 2. Run server:

        $ pip install bottle
        $ python server.py

 3. Browse http://localhost:8080/

TODO: Detect and report invalid UML inputs.
"""
from bottle import route, run, template, response, BaseResponse
from optparse import Values
from subprocess import check_output
from tempfile import SpooledTemporaryFile
import suml.common
import suml.yuml2dot

@route('/image/')
@route('/image/<spec:path>')
def image(spec=' '):
    spec = spec.replace('\n', ',')
    fout = SpooledTemporaryFile()

    # Execute Scruffy `suml`.
    options = Values(({
        'scruffy': True,
        'font': suml.common.defaultScruffyFont(),
        'png': True,
        'shadow': False,
    }))
    suml.yuml2dot.transform(spec, fout, options)

    fout.seek(0)
    png = fout.read()
    fout.close()

    # Server the generated image.
    response.content_type = 'image/png'
    return png

@route('/')
@route('/edit/')
@route('/edit/<uml:path>')
def index(uml='// Cool Class Diagram,[ICustomer|+name;+email|]^-[Customer],[Customer]<>-orders*>[Order],[Order]++-0..*>[LineItem],[Order]-[note:Aggregate root.]'):
    uml = uml.replace(',', '\n')
    return template("""
        <!DOCTYPE html>

            <title>Scruffy</title>

            <link rel="stylesheet" href="//code.jquery.com/ui/1.10.4/themes/smoothness/jquery-ui.css">

            <form>
                <textarea name="uml" rows="10" cols="80" autofocus="autofocus">{{uml}}</textarea>

See Scruffy syntax.

            </form>

        """,
        uml=uml)

if __name__ == "__main__":
    run(host='0.0.0.0', port=8080)
<commit_msg>Only hide when the UML was successfully generated.<commit_after>"""
 1. Install Scruffy
 2. Run server:

        $ pip install bottle
        $ python server.py

 3. Browse http://localhost:8080/

TODO: Detect and report invalid UML inputs.
"""
from bottle import route, run, template, response, BaseResponse
from optparse import Values
from subprocess import check_output
from tempfile import SpooledTemporaryFile
import suml.common
import suml.yuml2dot

@route('/image/')
@route('/image/<spec:path>')
def image(spec=' '):
    spec = spec.replace('\n', ',')
    fout = SpooledTemporaryFile()

    # Execute Scruffy `suml`.
    options = Values(({
        'scruffy': True,
        'font': suml.common.defaultScruffyFont(),
        'png': True,
        'shadow': False,
    }))
    suml.yuml2dot.transform(spec, fout, options)

    fout.seek(0)
    png = fout.read()
    fout.close()

    # Server the generated image.
    response.content_type = 'image/png'
    return png

@route('/')
@route('/edit/')
@route('/edit/<uml:path>')
def index(uml='// Cool Class Diagram,[ICustomer|+name;+email|]^-[Customer],[Customer]<>-orders*>[Order],[Order]++-0..*>[LineItem],[Order]-[note:Aggregate root.]'):
    uml = uml.replace(',', '\n')
    return template("""
        <!DOCTYPE html>

            <title>Scruffy</title>

            <link rel="stylesheet" href="//code.jquery.com/ui/1.10.4/themes/smoothness/jquery-ui.css">

            <form>
                <textarea name="uml" rows="10" cols="80" autofocus="autofocus">{{uml}}</textarea>

See Scruffy syntax.

            </form>

        """,
        uml=uml)

if __name__ == "__main__":
    run(host='0.0.0.0', port=8080)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import numpy as np
from abel.tools.polar import index_coords, cart2polar
import scipy.constants as const

# This file includes functions that have a known analytical Abel transform.
# They are used in unit testing as well for comparing different Abel impementations.
# See BaseAnalytical class for more information.

def sample_image_dribinski(n=361):
    """
    Sample test image used in the BASEX paper Rev. Sci. Instrum. 73, 2634 (2002)
    9x Gaussian functions of half-width 4 pixel + 1 background width 3600
      - anisotropy - ß = -1  for cosθ term
                   - ß = +2  for sinθ term
                   - ß =  0  isotropic, no angular variation
    (there are some missing negative exponents in the publication)

    Parameters
    ----------
    n: integer
      image size n rows x n cols

    Returns
    -------
    IM: 2D np.array
      image

    """

    def Gauss (r, r0, sigma2):
        return np.exp(-(r-r0)**2/sigma2)

    def intensity(r, theta):  # intensity function Eq. (16)
        sinetheta2 = np.sin(theta)**2
        cosinetheta2 = np.cos(theta)**2

        t0 = 7*Gauss(r, 10, 4)*sinetheta2
        t1 = 3*Gauss(r, 15, 4)
        t2 = 5*Gauss(r, 20, 4)*cosinetheta2

        t3 =   Gauss(r, 70, 4)
        t4 = 2*Gauss(r, 85, 4)*cosinetheta2
        t5 =   Gauss(r, 100, 4)*sinetheta2

        t6 = 2*Gauss(r, 145, 4)*sinetheta2
        t7 =   Gauss(r, 150, 4)
        t8 = 3*Gauss(r, 155, 4)*cosinetheta2

        t9 = 20*Gauss(r, 45, 3600)  # background under t3 to t5

        return 2000*(t0+t1+t2) + 200*(t3+t4+t5) + 50*(t6+t7+t8) + t9

    # meshgrid @DanHickstein issue #67, updated #70

    x = np.linspace(-180, 180, n)
    X, Y = np.meshgrid(x, x)
    R, THETA = cart2polar(X, Y)
    IM = intensity(R, THETA)

    return IM

def sample_image_O(n=1001, sigma=3, temperature=200):
    """
    Simulate the photoelectron spectrum of O- photodetachment
    3PJ <- 2P3/2,1/2

    Parameters
    ----------
    n: integer
        image size n rows x n columns

    sigma: float
        Gaussian 1/e width (pixels)

    temperature: float
        anion levels have Boltzmann population weight
        (2J+1) exp(-177.1 h c 100/k/temperature)

    """
    n2 = n//2 + n % 2

    x = np.linspace(-n2, n2, n)
    X, Y = np.meshgrid(x, x)

    # 2D Gaussian intensity distribution
    r = np.sqrt(X*X + Y*Y)

    # 3P2 <- 2P3/2
    IM = np.exp(-(r-341)**2/sigma**2)
    # 3P1 <- 2P3/2
    IM += 0.8*np.exp(-(r-285)**2/sigma**2)
    # 3P0 <- 2P3/2
    IM += 0.36*np.exp(-(r-257)**2/sigma**2)

    boltzmann = 0.5*np.exp(-177.1*const.h*const.c*100/const.k/temperature)
    # 3P2 <- 2P1/2
    IM += 0.2*boltzmann*np.exp(-(r-394)**2/sigma**2)
    # 3P1 <- 2P1/2
    IM += 0.36*boltzmann*np.exp(-(r-348)**2/sigma**2)
    # 3P0 <- 2P1/2
    IM += 0.16*boltzmann*np.exp(-(r-324)**2/sigma**2)

    return IM

class BaseAnalytical(object):

    def __init__(self, n, r_max, symmetric=True, **args):
        """This is the base class for functions that have a known Abel transform.
        Every such class should expose the following public attributes:
          - self.r: vector of positions along the r axis
          - self.func: the values of the original function (same shape as self.r)
          - self.abel: the values of the Abel transform (same shape as self.r)
          - self.mask_valid: mask of the r interval where the function is well smoothed/well behaved
                           (no known artefacts in the inverse Abel reconstuction), typically excluding
                           the origin, the domain boundaries, and function discontinuities, that can
                           be used for unit testing.

        See GaussianAnalytical for a concrete example.

        Parameters
        ----------
           - n : int: number of points along the r axis
           - r_max: float: maximum r interval
          - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
        """
        self.n = n
        self.r_max = r_max

        assert r_max > 0

        if symmetric:
            self.r = np.linspace(-r_max, r_max, n)
        else:
            self.r = np.linspace(0, r_max, n)

        self.dr =  np.diff(self.r)[0]

class StepAnalytical(BaseAnalytical):
    def __init__(self, n, r_max, r1, r2, A0=1.0,
                                                ratio_valid_step=1.0, symmetric=True):
        """
        Define a a symmetric step function and calculate it's analytical
        Abel transform. See examples/example_step.py

        Parameters
           - n : int: number of points along the r axis
           - r_max: float: range of the symmetric r interval
           - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
           - r1, r2: floats: bounds of the step function if r > 0
                    ( symetic function is constructed for r < 0)
           - A0: float: height of the step
           - ratio_valid_step: float: in the benchmark take only the central ratio*100% of the step
                                         (exclude possible artefacts on the edges)

        see https://github.com/PyAbel/PyAbel/pull/16

        """

        super(StepAnalytical, self).__init__(n, r_max, symmetric)

        self.r1, self.r2 = r1, r2
        self.A0 = A0

        mask = np.abs(np.abs(self.r)- 0.5*(r1 + r2)) < 0.5*(r2 - r1)

        fr = np.zeros(self.r.shape)
        fr[mask] = A0

        self.func = fr

        self.abel = sym_abel_step_1d(self.r, self.A0, self.r1, self.r2)[0]

        # exclude the region near the discontinuity
        self.mask_valid = np.abs(np.abs(self.r)- 0.5*(r1 + r2)) < \
                                                ratio_valid_step*0.5*(r2 - r1)

class GaussianAnalytical(BaseAnalytical):
    def __init__(self, n, r_max, sigma=1.0, A0=1.0,
                            ratio_valid_sigma=2.0, symmetric=True):
        """
        Define a gaussian function and calculate its analytical
        Abel transform. See examples/example_gaussian.py

        Parameters
           - n : int: number of points along the r axis
           - r_max: float: range of the symmetric r interval
           - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
           - sigma: floats: sigma parameter for the gaussian
           - A0: float: amplitude of the gaussian
           - ratio_valid_sigma: float: in the benchmark ta
                        0 < r < ration_valid_sigma * sigma
                        (exclude possible artefacts on the axis, and )

        Source: http://mathworld.wolfram.com/AbelTransform.html
        """

        super(GaussianAnalytical, self).__init__(n, r_max, symmetric)

        self.sigma = sigma
        self.A0 = A0

        r = self.r

        self.func = A0*np.exp(-r**2/sigma**2)

        self.abel = sigma * np.sqrt(np.pi) * A0*np.exp(-r**2/sigma**2)

        # exclude the region near the discontinuity
        self.mask_valid = (np.abs(self.r) < ratio_valid_sigma*sigma) &\
                          (np.abs(self.r) > 0)

def abel_step_analytical(r, A0, r0, r1):
    """
    Direct Abel transform of a step function located between r0 and r1,
    with a height A0

     A0 +                  +-------------+
        |                  |             |
        |                  |             |
      0 | -----------------+             +-------------
        +------------------+-------------+------------>
        0                  r0            r1           r axis

    This function is mostly used for unit testing the inverse Abel transform

    Parameters:

       r:   1D array, vecor of positions along the r axis. Must start with 0.
       r0, r1: floats, positions of the step along the r axis
       A0:  float or 1D array: height of the step. If 1D array the height can be
            variable along the Z axis

    Returns:
       1D array if A0 is a float, a 2D array otherwise
    """

    if np.all(r[[0,-1]]) :
        raise ValueError('The vector of r coordinates must start with 0.0')

    F_1d = np.zeros(r.shape)
    mask = (r>=r0)*(r<r1)
    F_1d[mask] = 2*np.sqrt(r1**2 - r[mask]**2)
    mask = r<r0
    F_1d[mask] = 2*np.sqrt(r1**2 - r[mask]**2) - 2*np.sqrt(r0**2 - r[mask]**2)
    A0 = np.atleast_1d(A0)
    if A0.ndim == 1:
        A0 = A0[:,np.newaxis]
    return F_1d[np.newaxis,:]*A0

def sym_abel_step_1d(r, A0, r0, r1):
    """
    Produces a symmetrical analytical transform of a 1d step
    """
    A0 = np.atleast_1d(A0)
    d = np.empty(A0.shape + r.shape)
    A0 = A0[np.newaxis, :]
    for sens, mask in enumerate([r>=0, r<=0]):
        d[:,mask] =  abel_step_analytical(np.abs(r[mask]), A0, r0, r1)

    return d
<commit_msg>sample_image_O() scaling factor to maintain peak positions, independent of image size<commit_after># -*- coding: utf-8 -*-
import numpy as np
from abel.tools.polar import index_coords, cart2polar
import scipy.constants as const

# This file includes functions that have a known analytical Abel transform.
# They are used in unit testing as well for comparing different Abel impementations.
# See BaseAnalytical class for more information.

def sample_image_dribinski(n=361):
    """
    Sample test image used in the BASEX paper Rev. Sci. Instrum. 73, 2634 (2002)
    9x Gaussian functions of half-width 4 pixel + 1 background width 3600
      - anisotropy - ß = -1  for cosθ term
                   - ß = +2  for sinθ term
                   - ß =  0  isotropic, no angular variation
    (there are some missing negative exponents in the publication)

    Parameters
    ----------
    n: integer
      image size n rows x n cols

    Returns
    -------
    IM: 2D np.array
      image

    """

    def Gauss (r, r0, sigma2):
        return np.exp(-(r-r0)**2/sigma2)

    def intensity(r, theta):  # intensity function Eq. (16)
        sinetheta2 = np.sin(theta)**2
        cosinetheta2 = np.cos(theta)**2

        t0 = 7*Gauss(r, 10, 4)*sinetheta2
        t1 = 3*Gauss(r, 15, 4)
        t2 = 5*Gauss(r, 20, 4)*cosinetheta2

        t3 =   Gauss(r, 70, 4)
        t4 = 2*Gauss(r, 85, 4)*cosinetheta2
        t5 =   Gauss(r, 100, 4)*sinetheta2

        t6 = 2*Gauss(r, 145, 4)*sinetheta2
        t7 =   Gauss(r, 150, 4)
        t8 = 3*Gauss(r, 155, 4)*cosinetheta2

        t9 = 20*Gauss(r, 45, 3600)  # background under t3 to t5

        return 2000*(t0+t1+t2) + 200*(t3+t4+t5) + 50*(t6+t7+t8) + t9

    # meshgrid @DanHickstein issue #67, updated #70

    x = np.linspace(-180, 180, n)
    X, Y = np.meshgrid(x, x)
    R, THETA = cart2polar(X, Y)
    IM = intensity(R, THETA)

    return IM

def sample_image_O(n=1001, sigma=3, temperature=200):
    """
    Simulate the photoelectron spectrum of O- photodetachment
    3PJ <- 2P3/2,1/2

    Parameters
    ----------
    n: integer
        image size n rows x n columns

    sigma: float
        Gaussian 1/e width (pixels)

    temperature: float
        anion levels have Boltzmann population weight
        (2J+1) exp(-177.1 h c 100/k/temperature)

    """
    n2 = n//2 + n % 2

    x = np.linspace(-n2, n2, n)
    X, Y = np.meshgrid(x, x)

    # 2D Gaussian intensity distribution
    r = np.sqrt(X*X + Y*Y)

    rfact = n/1001.0

    # 3P2 <- 2P3/2
    IM = np.exp(-(r-341*rfact)**2/sigma**2)
    # 3P1 <- 2P3/2
    IM += 0.8*np.exp(-(r-285*rfact)**2/sigma**2)
    # 3P0 <- 2P3/2
    IM += 0.36*np.exp(-(r-257*rfact)**2/sigma**2)

    boltzmann = 0.5*np.exp(-177.1*const.h*const.c*100/const.k/temperature)
    # 3P2 <- 2P1/2
    IM += 0.2*boltzmann*np.exp(-(r-394*rfact)**2/sigma**2)
    # 3P1 <- 2P1/2
    IM += 0.36*boltzmann*np.exp(-(r-348*rfact)**2/sigma**2)
    # 3P0 <- 2P1/2
    IM += 0.16*boltzmann*np.exp(-(r-324*rfact)**2/sigma**2)

    return IM

class BaseAnalytical(object):

    def __init__(self, n, r_max, symmetric=True, **args):
        """This is the base class for functions that have a known Abel transform.
        Every such class should expose the following public attributes:
          - self.r: vector of positions along the r axis
          - self.func: the values of the original function (same shape as self.r)
          - self.abel: the values of the Abel transform (same shape as self.r)
          - self.mask_valid: mask of the r interval where the function is well smoothed/well behaved
                           (no known artefacts in the inverse Abel reconstuction), typically excluding
                           the origin, the domain boundaries, and function discontinuities, that can
                           be used for unit testing.

        See GaussianAnalytical for a concrete example.

        Parameters
        ----------
           - n : int: number of points along the r axis
           - r_max: float: maximum r interval
          - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
        """
        self.n = n
        self.r_max = r_max

        assert r_max > 0

        if symmetric:
            self.r = np.linspace(-r_max, r_max, n)
        else:
            self.r = np.linspace(0, r_max, n)

        self.dr =  np.diff(self.r)[0]

class StepAnalytical(BaseAnalytical):
    def __init__(self, n, r_max, r1, r2, A0=1.0,
                                                ratio_valid_step=1.0, symmetric=True):
        """
        Define a a symmetric step function and calculate it's analytical
        Abel transform. See examples/example_step.py

        Parameters
           - n : int: number of points along the r axis
           - r_max: float: range of the symmetric r interval
           - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
           - r1, r2: floats: bounds of the step function if r > 0
                    ( symetic function is constructed for r < 0)
           - A0: float: height of the step
           - ratio_valid_step: float: in the benchmark take only the central ratio*100% of the step
                                         (exclude possible artefacts on the edges)

        see https://github.com/PyAbel/PyAbel/pull/16

        """

        super(StepAnalytical, self).__init__(n, r_max, symmetric)

        self.r1, self.r2 = r1, r2
        self.A0 = A0

        mask = np.abs(np.abs(self.r)- 0.5*(r1 + r2)) < 0.5*(r2 - r1)

        fr = np.zeros(self.r.shape)
        fr[mask] = A0

        self.func = fr

        self.abel = sym_abel_step_1d(self.r, self.A0, self.r1, self.r2)[0]

        # exclude the region near the discontinuity
        self.mask_valid = np.abs(np.abs(self.r)- 0.5*(r1 + r2)) < \
                                                ratio_valid_step*0.5*(r2 - r1)

class GaussianAnalytical(BaseAnalytical):
    def __init__(self, n, r_max, sigma=1.0, A0=1.0,
                            ratio_valid_sigma=2.0, symmetric=True):
        """
        Define a gaussian function and calculate its analytical
        Abel transform. See examples/example_gaussian.py

        Parameters
           - n : int: number of points along the r axis
           - r_max: float: range of the symmetric r interval
           - symmetric: if True the r interval is [-r_max, r_max]  (and n should be odd)
                       otherwise the r interval is [0, r_max]
           - sigma: floats: sigma parameter for the gaussian
           - A0: float: amplitude of the gaussian
           - ratio_valid_sigma: float: in the benchmark ta
                        0 < r < ration_valid_sigma * sigma
                        (exclude possible artefacts on the axis, and )

        Source: http://mathworld.wolfram.com/AbelTransform.html
        """

        super(GaussianAnalytical, self).__init__(n, r_max, symmetric)

        self.sigma = sigma
        self.A0 = A0

        r = self.r

        self.func = A0*np.exp(-r**2/sigma**2)

        self.abel = sigma * np.sqrt(np.pi) * A0*np.exp(-r**2/sigma**2)

        # exclude the region near the discontinuity
        self.mask_valid = (np.abs(self.r) < ratio_valid_sigma*sigma) &\
                          (np.abs(self.r) > 0)

def abel_step_analytical(r, A0, r0, r1):
    """
    Direct Abel transform of a step function located between r0 and r1,
    with a height A0

     A0 +                  +-------------+
        |                  |             |
        |                  |             |
      0 | -----------------+             +-------------
        +------------------+-------------+------------>
        0                  r0            r1           r axis

    This function is mostly used for unit testing the inverse Abel transform

    Parameters:

       r:   1D array, vecor of positions along the r axis. Must start with 0.
       r0, r1: floats, positions of the step along the r axis
       A0:  float or 1D array: height of the step. If 1D array the height can be
            variable along the Z axis

    Returns:
       1D array if A0 is a float, a 2D array otherwise
    """

    if np.all(r[[0,-1]]) :
        raise ValueError('The vector of r coordinates must start with 0.0')

    F_1d = np.zeros(r.shape)
    mask = (r>=r0)*(r<r1)
    F_1d[mask] = 2*np.sqrt(r1**2 - r[mask]**2)
    mask = r<r0
    F_1d[mask] = 2*np.sqrt(r1**2 - r[mask]**2) - 2*np.sqrt(r0**2 - r[mask]**2)
    A0 = np.atleast_1d(A0)
    if A0.ndim == 1:
        A0 = A0[:,np.newaxis]
    return F_1d[np.newaxis,:]*A0

def sym_abel_step_1d(r, A0, r0, r1):
    """
    Produces a symmetrical analytical transform of a 1d step
    """
    A0 = np.atleast_1d(A0)
    d = np.empty(A0.shape + r.shape)
    A0 = A0[np.newaxis, :]
    for sens, mask in enumerate([r>=0, r<=0]):
        d[:,mask] =  abel_step_analytical(np.abs(r[mask]), A0, r0, r1)

    return d
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase

from .. import models
from .. import factories

class GameTests(APITestCase):
    def test_create_game(self):
        """Ensure that we can create a game."""
        url = reverse('game-list')
        data = {'players': [], 'companies': []}
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create game: " + str(response.data))
        self.assertEqual(models.Game.objects.count(), 1)

class PlayerTests(APITestCase):
    def test_create_player(self):
        """Ensure that we can create players."""
        game = factories.GameFactory.create()
        url = reverse('player-list')
        data = {'name': 'Alice', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a player: " + str(response.data))
        self.assertEqual(models.Player.objects.count(), 1)

    def test_cannot_create_duplicate_player_for_single_game(self):
        """Disallow creating two players with the same name in a game."""
        game = factories.GameFactory.create()
        player = factories.PlayerFactory.create(game=game, name='Alice')
        url = reverse('player-list')
        data = {'name': 'Alice', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate player: " + str(response.data))
        self.assertEqual(models.Player.objects.count(), 1)

class CompanyTests(APITestCase):
    def test_create_company(self):
        """Ensure that we can create companies."""
        game = factories.GameFactory.create()
        url = reverse('company-list')
        data = {'name': 'B&O', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a company: " + str(response.data))
        self.assertEqual(models.Company.objects.count(), 1)

    def test_cannot_create_duplicate_company_for_single_game(self):
        """Disallow creating two companies with the same name in a game."""
        game = factories.GameFactory.create()
        company = factories.CompanyFactory.create(game=game, name='B&O')
        url = reverse('company-list')
        data = {'name': 'B&O', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate company: " + str(response.data))
        self.assertEqual(models.Company.objects.count(), 1)

class PlayerShareTests(APITestCase):
    def test_create_share(self):
        """Ensure that we can create shares."""
        game = factories.GameFactory.create()
        player = factories.PlayerFactory.create(game=game)
        company = factories.CompanyFactory.create(game=game)
        url = reverse('playershare-list')
        data = {'game': game.pk, 'owner': player.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a share: " + str(response.data))
        self.assertEqual(models.PlayerShare.objects.count(), 1)

    def test_cannot_create_duplicate_share_holdings(self):
        """
        Ensure that a player doesn't have two share holding records for
        a single company
        """
        game = factories.GameFactory.create()
        player = factories.PlayerFactory.create(game=game)
        company = factories.CompanyFactory.create(game=game)
        share = factories.PlayerShareFactory.create(owner=player,
            company=company)
        url = reverse('playershare-list')
        data = {'owner': player.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate share holdings: " + str(response.data))
        self.assertEqual(models.PlayerShare.objects.count(), 1)
        self.assertIn('non_field_errors', response.data.keys())

class CompanyShareTests(APITestCase):
    def test_create_self_owning_share(self):
        """Ensure that we can create company shares."""
        game = factories.GameFactory.create()
        company = factories.CompanyFactory.create(game=game)
        url = reverse('companyshare-list')
        data = {'owner': company.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a self owning share: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)

    def test_create_share_owning_other_company(self):
        """Ensure that companies can own shares in other companies."""
        game = factories.GameFactory.create()
        company1, company2 = factories.CompanyFactory.create_batch(size=2,
            game=game)
        url = reverse('companyshare-list')
        data = {'owner': company1.pk, 'company': company2.pk}

        response = self.client.post(url, data, format='json')

        share = models.CompanyShare.objects.first()
        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a company share: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)
        self.assertEqual(share.owner, company1)
        self.assertEqual(share.company, company2)

    def test_cannot_create_duplicate_share_holdings(self):
        """
        Ensure that a company doesn't have two share holding records for a
        single company
        """
        game = factories.GameFactory.create()
        company1, company2 = factories.CompanyFactory.create_batch(size=2,
            game=game)
        share = factories.CompanyShareFactory.create(owner=company1,
            company=company2)
        url = reverse('companyshare-list')
        data = {'owner': company1.pk, 'company': company2.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate share holdings: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)
        self.assertIn('non_field_errors', response.data.keys())
<commit_msg>Create one game for API tests to speed them up<commit_after># -*- coding: utf-8 -*-
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase

from .. import models
from .. import factories

game = None

def setUpModule():
    global game
    game = factories.GameFactory.create()

def tearDownModule():
    game.delete()

class GameTests(APITestCase):
    def test_create_game(self):
        """Ensure that we can create a game."""
        url = reverse('game-list')
        data = {'players': [], 'companies': []}
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create game: " + str(response.data))
        # We have 2 games since there is also a global game
        self.assertEqual(models.Game.objects.count(), 2)

class PlayerTests(APITestCase):
    def test_create_player(self):
        """Ensure that we can create players."""
        url = reverse('player-list')
        data = {'name': 'Alice', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a player: " + str(response.data))
        self.assertEqual(models.Player.objects.count(), 1)

    def test_cannot_create_duplicate_player_for_single_game(self):
        """Disallow creating two players with the same name in a game."""
        player = factories.PlayerFactory.create(game=game, name='Alice')
        url = reverse('player-list')
        data = {'name': 'Alice', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate player: " + str(response.data))
        self.assertEqual(models.Player.objects.count(), 1)

class CompanyTests(APITestCase):
    def test_create_company(self):
        """Ensure that we can create companies."""
        url = reverse('company-list')
        data = {'name': 'B&O', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a company: " + str(response.data))
        self.assertEqual(models.Company.objects.count(), 1)

    def test_cannot_create_duplicate_company_for_single_game(self):
        """Disallow creating two companies with the same name in a game."""
        company = factories.CompanyFactory.create(game=game, name='B&O')
        url = reverse('company-list')
        data = {'name': 'B&O', 'game': game.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate company: " + str(response.data))
        self.assertEqual(models.Company.objects.count(), 1)

class PlayerShareTests(APITestCase):
    def test_create_share(self):
        """Ensure that we can create shares."""
        player = factories.PlayerFactory.create(game=game)
        company = factories.CompanyFactory.create(game=game)
        url = reverse('playershare-list')
        data = {'game': game.pk, 'owner': player.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a share: " + str(response.data))
        self.assertEqual(models.PlayerShare.objects.count(), 1)

    def test_cannot_create_duplicate_share_holdings(self):
        """
        Ensure that a player doesn't have two share holding records for
        a single company
        """
        player = factories.PlayerFactory.create(game=game)
        company = factories.CompanyFactory.create(game=game)
        share = factories.PlayerShareFactory.create(owner=player,
            company=company)
        url = reverse('playershare-list')
        data = {'owner': player.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate share holdings: " + str(response.data))
        self.assertEqual(models.PlayerShare.objects.count(), 1)
        self.assertIn('non_field_errors', response.data.keys())

class CompanyShareTests(APITestCase):
    def test_create_self_owning_share(self):
        """Ensure that we can create company shares."""
        company = factories.CompanyFactory.create(game=game)
        url = reverse('companyshare-list')
        data = {'owner': company.pk, 'company': company.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a self owning share: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)

    def test_create_share_owning_other_company(self):
        """Ensure that companies can own shares in other companies."""
        company1, company2 = factories.CompanyFactory.create_batch(size=2,
            game=game)
        url = reverse('companyshare-list')
        data = {'owner': company1.pk, 'company': company2.pk}

        response = self.client.post(url, data, format='json')

        share = models.CompanyShare.objects.first()
        self.assertEqual(response.status_code, status.HTTP_201_CREATED,
            "Could not create a company share: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)
        self.assertEqual(share.owner, company1)
        self.assertEqual(share.company, company2)

    def test_cannot_create_duplicate_share_holdings(self):
        """
        Ensure that a company doesn't have two share holding records for a
        single company
        """
        company1, company2 = factories.CompanyFactory.create_batch(size=2,
            game=game)
        share = factories.CompanyShareFactory.create(owner=company1,
            company=company2)
        url = reverse('companyshare-list')
        data = {'owner': company1.pk, 'company': company2.pk}

        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST,
            "Created duplicate share holdings: " + str(response.data))
        self.assertEqual(models.CompanyShare.objects.count(), 1)
        self.assertIn('non_field_errors', response.data.keys())
<|endoftext|>


<commit_before>import requests
import settings
import functools
from flask import g, redirect, url_for, request
from hashids import Hashids
import boto3
from itsdangerous import URLSafeTimedSerializer, URLSafeSerializer
import sendgrid
from sendgrid.helpers.mail import *
from premailer import transform
from datetime import date
from urlparse import urlparse, urljoin
import logging

resume_hash = Hashids(min_length=8, salt=settings.RESUME_HASH_SALT)
s3 = boto3.resource('s3', aws_access_key_id=settings.AWS_ACCESS_KEY,
                    aws_secret_access_key=settings.AWS_SECRET_KEY)
s3_client = boto3.client(
    's3', aws_access_key_id=settings.AWS_ACCESS_KEY, aws_secret_access_key=settings.AWS_SECRET_KEY)
timed_serializer = URLSafeTimedSerializer(settings.SECRET_KEY)
serializer = URLSafeSerializer(settings.SECRET_KEY)
sg = sendgrid.SendGridAPIClient(apikey=settings.SENDGRID_API_KEY)

def validate_email(email):
    return requests.get(
        "https://api.mailgun.net/v3/address/validate",
        auth=("api", settings.MAILGUN_PUB_KEY),
        params={"address": email})

# A user who has any of these roles will be permitted to view this page
def roles_required(*role_names):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if not g.user.is_authenticated:
                return redirect(url_for('corp-login'))
            user_roles = get_current_user_roles()
            authorized = False
            for role in role_names:
                if role in user_roles:
                    authorized = True
                    break
            if not authorized:
                return 'Unauthorized'
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def get_current_user_roles():
    return set(role.name for role in g.user.roles)

# Require user to be logged in and redirect unauthenticated users to corporate login screen
def corp_login_required(f):
    @functools.wraps(f)
    def decorated_view(*args, **kwargs):
        if not g.user.is_authenticated:
            return redirect(url_for('login'))
        roles = get_current_user_roles()
        if 'admin' not in roles and 'corp' not in roles:
            return redirect(url_for('login'))
        return f(*args, **kwargs)

    return decorated_view

def user_status_whitelist(*statuses):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if 'admin' not in get_current_user_roles() and g.user.status not in statuses:
                return redirect(url_for(get_default_dashboard_for_role()))
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def user_extra_application_required(func):
    @functools.wraps(func)
    def decorated_view(*args, **kwargs):
        if 'admin' not in get_current_user_roles():
            if not g.user.campus_ambassador and not g.user.needs_travel_reimbursement:
                return redirect(url_for(get_default_dashboard_for_role()))
        return func(*args, **kwargs)

    return decorated_view

def user_status_blacklist(*statuses):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if 'admin' not in get_current_user_roles() and g.user.status in statuses:
                return redirect(url_for(get_default_dashboard_for_role()))
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def get_default_dashboard_for_role():
    roles = get_current_user_roles()
    if "admin" in roles:
        return "admin-dash"
    if "corp" in roles:
        return "corp-dash"
    return "dashboard"

def redirect_to_dashboard_if_authed(func):
    @functools.wraps(func)
    def decorated_view(*args, **kwargs):
        if g.user.is_authenticated:
            return redirect(url_for(get_default_dashboard_for_role()))
        return func(*args, **kwargs)

    return decorated_view

def send_email(from_email, subject, to_email, txt_content=None, html_content=None, attachments=None, from_name=None):
    if not from_name:
        from_name = settings.HACKATHON_NAME

    if not from_email:
        logging.error("No from_email found while sending emails")
        raise RuntimeError("No from_email")

    sg_mail = Mail()
    sg_mail.from_email = Email(from_email, from_name)
    personalization = Personalization()
    personalization.add_to(Email(to_email))
    personalization.subject = subject
    sg_mail.add_personalization(personalization)
    transform(html_content)
    if txt_content is None and html_content is None:
        raise ValueError('Must send some type of content')
    if txt_content:
        sg_mail.add_content(Content('text/plain', txt_content))
    if html_content:
        sg_mail.add_content(Content('text/html', transform(html_content)))

    if attachments:
        for attachment in attachments:
            print(str(attachment))
            sg_attachment = Attachment()
            sg_attachment['content'] = attachment['encoded']
            sg_attachment['filename'] = attachment['filename']
            sg_attachment['type'] = attachment['file_type']
            sg_attachment['disposition'] = 'attachment'
            sg_attachment['content_id'] = attachment['encoded']
            sg_mail.add_attachment(sg_attachment)

    mail_body = sg_mail.get()
    try:
        response = sg.client.mail.send.post(request_body=mail_body)
    except Exception as e:
        logging.error("Error mailing using SendGrid, message: {}".format(e.message))
        logging.error("error body: {}".format(e.body))
        logging.error("email: {}".format(mail_body))
        raise e
    if response.status_code != 202:
        raise RuntimeError  # TODO: Get a proper error on this

def calculate_age(dob):
    today = date.today()
    return today.year - dob.year - ((today.month, today.day) < (dob.month, dob.day))

def is_safe_url(target):
    ref_url = urlparse(request.host_url)
    test_url = urlparse(urljoin(request.host_url, target))
    return test_url.scheme in ('http', 'https') and (ref_url.netloc == test_url.netloc)
<commit_msg>sendgrid doesn't allow you set attributes<commit_after>import requests
import settings
import functools
from flask import g, redirect, url_for, request
from hashids import Hashids
import boto3
from itsdangerous import URLSafeTimedSerializer, URLSafeSerializer
import sendgrid
from sendgrid.helpers.mail import *
from premailer import transform
from datetime import date
from urlparse import urlparse, urljoin
import logging

resume_hash = Hashids(min_length=8, salt=settings.RESUME_HASH_SALT)
s3 = boto3.resource('s3', aws_access_key_id=settings.AWS_ACCESS_KEY,
                    aws_secret_access_key=settings.AWS_SECRET_KEY)
s3_client = boto3.client(
    's3', aws_access_key_id=settings.AWS_ACCESS_KEY, aws_secret_access_key=settings.AWS_SECRET_KEY)
timed_serializer = URLSafeTimedSerializer(settings.SECRET_KEY)
serializer = URLSafeSerializer(settings.SECRET_KEY)
sg = sendgrid.SendGridAPIClient(apikey=settings.SENDGRID_API_KEY)

def validate_email(email):
    return requests.get(
        "https://api.mailgun.net/v3/address/validate",
        auth=("api", settings.MAILGUN_PUB_KEY),
        params={"address": email})

# A user who has any of these roles will be permitted to view this page
def roles_required(*role_names):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if not g.user.is_authenticated:
                return redirect(url_for('corp-login'))
            user_roles = get_current_user_roles()
            authorized = False
            for role in role_names:
                if role in user_roles:
                    authorized = True
                    break
            if not authorized:
                return 'Unauthorized'
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def get_current_user_roles():
    return set(role.name for role in g.user.roles)

# Require user to be logged in and redirect unauthenticated users to corporate login screen
def corp_login_required(f):
    @functools.wraps(f)
    def decorated_view(*args, **kwargs):
        if not g.user.is_authenticated:
            return redirect(url_for('login'))
        roles = get_current_user_roles()
        if 'admin' not in roles and 'corp' not in roles:
            return redirect(url_for('login'))
        return f(*args, **kwargs)

    return decorated_view

def user_status_whitelist(*statuses):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if 'admin' not in get_current_user_roles() and g.user.status not in statuses:
                return redirect(url_for(get_default_dashboard_for_role()))
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def user_extra_application_required(func):
    @functools.wraps(func)
    def decorated_view(*args, **kwargs):
        if 'admin' not in get_current_user_roles():
            if not g.user.campus_ambassador and not g.user.needs_travel_reimbursement:
                return redirect(url_for(get_default_dashboard_for_role()))
        return func(*args, **kwargs)

    return decorated_view

def user_status_blacklist(*statuses):
    def wrapper(func):
        @functools.wraps(func)
        def decorated_view(*args, **kwargs):
            if 'admin' not in get_current_user_roles() and g.user.status in statuses:
                return redirect(url_for(get_default_dashboard_for_role()))
            return func(*args, **kwargs)

        return decorated_view

    return wrapper

def get_default_dashboard_for_role():
    roles = get_current_user_roles()
    if "admin" in roles:
        return "admin-dash"
    if "corp" in roles:
        return "corp-dash"
    return "dashboard"

def redirect_to_dashboard_if_authed(func):
    @functools.wraps(func)
    def decorated_view(*args, **kwargs):
        if g.user.is_authenticated:
            return redirect(url_for(get_default_dashboard_for_role()))
        return func(*args, **kwargs)

    return decorated_view

def send_email(from_email, subject, to_email, txt_content=None, html_content=None, attachments=None, from_name=None):
    if not from_name:
        from_name = settings.HACKATHON_NAME

    if not from_email:
        logging.error("No from_email found while sending emails")
        raise RuntimeError("No from_email")

    sg_mail = Mail()
    sg_mail.from_email = Email(from_email, from_name)
    personalization = Personalization()
    personalization.add_to(Email(to_email))
    personalization.subject = subject
    sg_mail.add_personalization(personalization)
    transform(html_content)
    if txt_content is None and html_content is None:
        raise ValueError('Must send some type of content')
    if txt_content:
        sg_mail.add_content(Content('text/plain', txt_content))
    if html_content:
        sg_mail.add_content(Content('text/html', transform(html_content)))

    if attachments:
        for attachment in attachments:
            print(str(attachment))
            sg_attachment = Attachment()
            sg_attachment.content = attachment['encoded']
            sg_attachment.filename = attachment['filename']
            sg_attachment.type = attachment['file_type']
            sg_attachment.disposition = 'attachment'
            sg_attachment.content_id = attachment['encoded']
            sg_mail.add_attachment(sg_attachment)

    mail_body = sg_mail.get()
    try:
        response = sg.client.mail.send.post(request_body=mail_body)
    except Exception as e:
        logging.error("Error mailing using SendGrid, message: {}".format(e.message))
        logging.error("error body: {}".format(e.body))
        logging.error("email: {}".format(mail_body))
        raise e
    if response.status_code != 202:
        raise RuntimeError  # TODO: Get a proper error on this

def calculate_age(dob):
    today = date.today()
    return today.year - dob.year - ((today.month, today.day) < (dob.month, dob.day))

def is_safe_url(target):
    ref_url = urlparse(request.host_url)
    test_url = urlparse(urljoin(request.host_url, target))
    return test_url.scheme in ('http', 'https') and (ref_url.netloc == test_url.netloc)
<|endoftext|>