<commit_before>import imgaug as ia
from imgaug import augmenters as iaa
import os
import cv2
import numpy as np
import logging
import time

OUTPUT_DIR = '/tmp/export_alpr_plate_plugin/'
NUM_AUG_IMAGES = 10
BBOX_SCALES = [ 0.75, 1, 1.25]
IM_HEIGHT = 64
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'

# Sometimes(0.5, ...) applies the given augmenter in 50% of all cases,
# e.g. Sometimes(0.5, GaussianBlur(0.3)) would blur roughly every second image.
sometimes = lambda aug: iaa.Sometimes(0.5, aug)

# execute 0 to 5 of the following (less important) augmenters per image
# don't execute all of them, as that would often be way too strong
im_filters = iaa.SomeOf((0, 5),
    [
        iaa.OneOf([
            iaa.GaussianBlur((0, 1.5)), # blur images with a sigma between 0 and 3.0
            iaa.AverageBlur(k=(2, 3)), # blur image using local means with kernel sizes between 2 and 7
            iaa.MedianBlur(k=(3, 3)), # blur image using local medians with kernel sizes between 2 and 7
        ]),
        iaa.Sharpen(alpha=(0, 0.1), lightness=(0.75, 1.5)), # sharpen images
        iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05*255), per_channel=0.5), # add gaussian noise to images
        iaa.Multiply((0.5, 1.5), per_channel=0.5), # change brightness of images (50-150% of original value)
        iaa.Add((-20, 20)), # change brightness of images (by -10 to 10 of original value)
        iaa.OneOf([
            iaa.ContrastNormalization((0.75, 1.75), per_channel=0.5), # improve or worsen the contrast
            iaa.ContrastNormalization((0.5, 1.5)), # improve or worsen the contrast
        ]),
        iaa.Grayscale(alpha=(0.0, 1.0)),
    ],
    random_order=True
)

# Define our sequence of augmentation steps that will be applied to every image
# All augmenters with per_channel=0.5 will sample one value _per image_
# in 50% of all cases. In all other cases they will sample new values
# _per channel_.
seq_3 = iaa.Sequential(
    [
        # apply the following augmenters to most images
        sometimes(iaa.Crop(percent=(0, 0.1))), # crop images by 0-10% of their height/width
        sometimes(iaa.Affine(
            scale={"x": (0.8, 1.2), "y": (0.8, 1.2)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, # translate by -20 to +20 percent (per axis)
            rotate=(-4, 4), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

seq_2 = iaa.Sequential(
    [
        # apply the following augmenters to most images
        sometimes(iaa.Crop(percent=(0, 0.05))), # crop images by 0-10% of their height/width
        sometimes(iaa.Affine(
            scale={"x": (0.9, 1.1), "y": (0.9, 1.1)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, # translate by -20 to +20 percent (per axis)
            rotate=(-3, 3), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

seq_1 = iaa.Sequential(
    [
        sometimes(iaa.Affine(
            scale={"x": (0.95, 1.05), "y": (0.95, 1.05)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.025, 0.025), "y": (-0.05, 0.025)}, # translate by -20 to +20 percent (per axis)
            rotate=(0, 0), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

class AnnotatorPlugin:
    _VERSION = '0.0.1'

    def __init__(self, dataset, partition):

        global OUTPUT_DIR
        OUTPUT_DIR += "/"+dataset['name']+"/"

        self.logger = logging.getLogger('annotator_supreme.plugin.alpr_export')
        self.logger.info("Exporting dataset %s", dataset['name'])
        self.image_path_list = [[],[]]
        self.label_list = [[],[]]

        if os.system('mkdir -p '+OUTPUT_DIR+'/train/') :
            self.logger.error("Cannot create %s/train/", OUTPUT_DIR)
            return

        if os.system('mkdir -p '+OUTPUT_DIR+'/val/') :
            self.logger.error("Cannot create %s/val/", OUTPUT_DIR)
            return

    def is_date(self, text):
        return len(text.split('/')) == 3

    def process(self, im, anno):
        '''
        main function of the plugin

        parameters:
            - im : an opencv image matrix
            - anno: a json representing the annotation for this image
        '''
        part = anno['partition']
        self.logger.info("partition %s", part)
        if part == 0:
            im_name = OUTPUT_DIR+'/train/'+anno['phash']
        else:
            im_name = OUTPUT_DIR+'/val/'+anno['phash']
        width   = float(im.shape[1])
        height  = float(im.shape[0])
        curr_tag = ''
        # print('processing', anno['phash'])

        annotations = []
        imgs_vec    = []
        txt_vec     = []
        curr_im = 0
        for bb in anno['anno']:
            if bb['ignore']:
                continue

            if len(bb['labels']) <= 0:
                continue

           #  if bb['labels'][0] != 'plate' and bb['labels'][0] != 'plate1' and bb['labels'][0] != 'plate2' and bb['labels'][0] != 'plate_ant':
           #     continue

            curr_label = ''
            if len(bb['labels']) > 0:
                curr_label = bb['labels'][0].lower().strip()
            curr_label = curr_label.replace(' ','').replace('-','')
            bad_letter = False

            #skip annotations with blank labels
            if curr_label == '' :
                continue

            for a in curr_label:
                if a not in alphabet:
                    self.logger.debug('invalid char %s %s', a, anno['phash'])
                    bad_letter = True

            if bad_letter:
                continue

            l = int(bb['left'])
            t = int(bb['top'])
            r = int(bb['right'])
            b = int(bb['bottom'])

            # Save the original cropped plate
            org_im = np.copy(im[t:b,l:r,:])
            org_name = im_name+'_org'

            #save label file
            with open(org_name+'.txt', 'w') as f:
                f.write(curr_label+'\n')

            #save image file
            cv2.imwrite(org_name+'.png', org_im)
            self.image_path_list[part].append(org_name+'.png')
            self.label_list[part].append(curr_label)

            #Do not augment images from the validation partition
            if part == 1:
                continue;

            for curr_scale in BBOX_SCALES  :
                (nl, nt, nr, nb) = self.safe_scale_bbox(l, t, r, b, curr_scale, im.shape)
                # imgs_vec = np.asarray([np.copy(im[nt:nb,nl:nr,:])])
                aux_im = np.copy(im[nt:nb,nl:nr,:])

                if aux_im.shape[0] <= 0 or aux_im.shape[1] <= 0:
                    break
                scale_im = IM_HEIGHT/float(aux_im.shape[0])
                aux_im = cv2.resize(aux_im, (int(aux_im.shape[1]*scale_im), IM_HEIGHT))
                imgs_vec = np.asarray([aux_im])
                txt_vec.append(curr_label)

                for aug in range(0,NUM_AUG_IMAGES):
                    # Apply different augmentations depending on the scale of the bounding box
                    if curr_scale == BBOX_SCALES[0] :
                        images_aug = seq_1.augment_images(imgs_vec)
                    elif curr_scale == BBOX_SCALES[1]:
                        images_aug = seq_2.augment_images(imgs_vec)
                    else:
                        images_aug = seq_3.augment_images(imgs_vec)

                    for j in range(0,len(images_aug)):
                        curr_name = im_name+'_'+str(curr_im).zfill(4)
                        with open(curr_name+'.txt', 'w') as f:
                            f.write(curr_label+'\n')
                        cv2.imwrite(curr_name+'.png', images_aug[j,:])
                        self.image_path_list[part].append(curr_name+'.png')
                        self.label_list[part].append(curr_label)
                        curr_im += 1

        self.logger.info('Done image %s', im_name)

        return (im, anno)

    def end(self):

        EXEC_ID=str(int(time.time()))  # id used for saving files
        with open(OUTPUT_DIR+'/train/image_path_list_'+EXEC_ID+'.txt', 'w') as f:
            for p in self.image_path_list[0]:
                f.write(p+'\n')

        with open(OUTPUT_DIR+'/train/label_list_'+EXEC_ID+'.txt', 'w') as f:
            for l in self.label_list[0]:
                f.write(l+'\n')

        with open(OUTPUT_DIR+'/val/image_path_list_'+EXEC_ID+'.txt', 'w') as f:
            for p in self.image_path_list[1]:
                f.write(p+'\n')

        with open(OUTPUT_DIR+'/val/label_list_'+EXEC_ID+'.txt', 'w') as f:
            for l in self.label_list[1]:
                f.write(l+'\n')

        self.logger.info('Execution id: %s', EXEC_ID)

    def get_parameters(self):
        return {'parameters': []}

    def get_version(self):
        return self._VERSION

    def safe_scale_bbox(self, l, t, r, b, scale, shape):
        dw = (r-l)*(scale-1)/6.0
        dh = (b-t)*(scale-1)/2.0

        l = int(l-dw)
        t = int(t-dh)
        r = int(r+dw)
        b = int(b+dh)

        if l<0: l=0
        if t<0: t=0
        if b>=shape[0]: b=shape[0]-1
        if r>=shape[1]: r=shape[1]-1

        return (l,t,r,b)
<commit_msg>FIX: alpr export plugin was overwriting cropped annotations of the same images<commit_after>import imgaug as ia
from imgaug import augmenters as iaa
import os
import cv2
import numpy as np
import logging
import time

OUTPUT_DIR = '/tmp/export_alpr_plate_plugin/'
NUM_AUG_IMAGES = 10
BBOX_SCALES = [ 0.75, 1, 1.25]
IM_HEIGHT = 64
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'

# Sometimes(0.5, ...) applies the given augmenter in 50% of all cases,
# e.g. Sometimes(0.5, GaussianBlur(0.3)) would blur roughly every second image.
sometimes = lambda aug: iaa.Sometimes(0.5, aug)

# execute 0 to 5 of the following (less important) augmenters per image
# don't execute all of them, as that would often be way too strong
im_filters = iaa.SomeOf((0, 5),
    [
        iaa.OneOf([
            iaa.GaussianBlur((0, 1.5)), # blur images with a sigma between 0 and 3.0
            iaa.AverageBlur(k=(2, 3)), # blur image using local means with kernel sizes between 2 and 7
            iaa.MedianBlur(k=(3, 3)), # blur image using local medians with kernel sizes between 2 and 7
        ]),
        iaa.Sharpen(alpha=(0, 0.1), lightness=(0.75, 1.5)), # sharpen images
        iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05*255), per_channel=0.5), # add gaussian noise to images
        iaa.Multiply((0.5, 1.5), per_channel=0.5), # change brightness of images (50-150% of original value)
        iaa.Add((-20, 20)), # change brightness of images (by -10 to 10 of original value)
        iaa.OneOf([
            iaa.ContrastNormalization((0.75, 1.75), per_channel=0.5), # improve or worsen the contrast
            iaa.ContrastNormalization((0.5, 1.5)), # improve or worsen the contrast
        ]),
        iaa.Grayscale(alpha=(0.0, 1.0)),
    ],
    random_order=True
)

# Define our sequence of augmentation steps that will be applied to every image
# All augmenters with per_channel=0.5 will sample one value _per image_
# in 50% of all cases. In all other cases they will sample new values
# _per channel_.
seq_3 = iaa.Sequential(
    [
        # apply the following augmenters to most images
        sometimes(iaa.Crop(percent=(0, 0.1))), # crop images by 0-10% of their height/width
        sometimes(iaa.Affine(
            scale={"x": (0.8, 1.2), "y": (0.8, 1.2)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, # translate by -20 to +20 percent (per axis)
            rotate=(-4, 4), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

seq_2 = iaa.Sequential(
    [
        # apply the following augmenters to most images
        sometimes(iaa.Crop(percent=(0, 0.05))), # crop images by 0-10% of their height/width
        sometimes(iaa.Affine(
            scale={"x": (0.9, 1.1), "y": (0.9, 1.1)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, # translate by -20 to +20 percent (per axis)
            rotate=(-3, 3), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

seq_1 = iaa.Sequential(
    [
        sometimes(iaa.Affine(
            scale={"x": (0.95, 1.05), "y": (0.95, 1.05)}, # scale images to 80-120% of their size, individually per axis
            translate_percent={"x": (-0.025, 0.025), "y": (-0.05, 0.025)}, # translate by -20 to +20 percent (per axis)
            rotate=(0, 0), # rotate by -45 to +45 degrees
            shear=(-6, 6), # shear by -16 to +16 degrees
            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
            mode='edge' # use any of scikit-image's warping modes (see 2nd image from the top for examples)
        )),
        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.1))),
        im_filters
    ],
    random_order=True
)

class AnnotatorPlugin:
    _VERSION = '0.0.1'

    def __init__(self, dataset, partition):

        global OUTPUT_DIR
        OUTPUT_DIR += "/"+dataset['name']+"/"

        self.logger = logging.getLogger('annotator_supreme.plugin.alpr_export')
        self.logger.info("Exporting dataset %s", dataset['name'])
        self.image_path_list = [[],[]]
        self.label_list = [[],[]]

        if os.system('mkdir -p '+OUTPUT_DIR+'/train/') :
            self.logger.error("Cannot create %s/train/", OUTPUT_DIR)
            return

        if os.system('mkdir -p '+OUTPUT_DIR+'/val/') :
            self.logger.error("Cannot create %s/val/", OUTPUT_DIR)
            return

    def is_date(self, text):
        return len(text.split('/')) == 3

    def process(self, im, anno):
        '''
        main function of the plugin

        parameters:
            - im : an opencv image matrix
            - anno: a json representing the annotation for this image
        '''
        part = anno['partition']
        self.logger.info("partition %s", part)
        if part == 0:
            im_name = OUTPUT_DIR+'/train/'+anno['phash']
        else:
            im_name = OUTPUT_DIR+'/val/'+anno['phash']
        width   = float(im.shape[1])
        height  = float(im.shape[0])
        curr_tag = ''
        # print('processing', anno['phash'])

        annotations = []
        imgs_vec    = []
        txt_vec     = []
        curr_im = 0
        for anno_id, bb in enumerate(anno['anno']):
            anno_name = im_name+'_'+str(anno_id)

            if bb['ignore']:
                continue

            if len(bb['labels']) <= 0:
                continue

           #  if bb['labels'][0] != 'plate' and bb['labels'][0] != 'plate1' and bb['labels'][0] != 'plate2' and bb['labels'][0] != 'plate_ant':
           #     continue

            curr_label = ''
            if len(bb['labels']) > 0:
                curr_label = bb['labels'][0].lower().strip()
            curr_label = curr_label.replace(' ','').replace('-','')
            bad_letter = False

            #skip annotations with blank labels
            if curr_label == '' :
                continue

            for a in curr_label:
                if a not in alphabet:
                    self.logger.debug('invalid char %s %s', a, anno['phash'])
                    bad_letter = True

            if bad_letter:
                continue

            l = int(bb['left'])
            t = int(bb['top'])
            r = int(bb['right'])
            b = int(bb['bottom'])

            # Save the original cropped plate
            org_im = np.copy(im[t:b,l:r,:])
            org_name = anno_name+'_org'

            #save label file
            with open(org_name+'.txt', 'w') as f:
                f.write(curr_label+'\n')

            #save image file
            cv2.imwrite(org_name+'.png', org_im)
            self.image_path_list[part].append(org_name+'.png')
            self.label_list[part].append(curr_label)

            #Do not augment images from the validation partition
            if part == 1:
                continue;

            for curr_scale in BBOX_SCALES  :
                (nl, nt, nr, nb) = self.safe_scale_bbox(l, t, r, b, curr_scale, im.shape)
                # imgs_vec = np.asarray([np.copy(im[nt:nb,nl:nr,:])])
                aux_im = np.copy(im[nt:nb,nl:nr,:])

                if aux_im.shape[0] <= 0 or aux_im.shape[1] <= 0:
                    break
                scale_im = IM_HEIGHT/float(aux_im.shape[0])
                aux_im = cv2.resize(aux_im, (int(aux_im.shape[1]*scale_im), IM_HEIGHT))
                imgs_vec = np.asarray([aux_im])
                txt_vec.append(curr_label)

                for aug in range(0,NUM_AUG_IMAGES):
                    # Apply different augmentations depending on the scale of the bounding box
                    if curr_scale == BBOX_SCALES[0] :
                        images_aug = seq_1.augment_images(imgs_vec)
                    elif curr_scale == BBOX_SCALES[1]:
                        images_aug = seq_2.augment_images(imgs_vec)
                    else:
                        images_aug = seq_3.augment_images(imgs_vec)

                    for j in range(0,len(images_aug)):
                        curr_name = anno_name+'_'+str(curr_im).zfill(4)
                        with open(curr_name+'.txt', 'w') as f:
                            f.write(curr_label+'\n')
                        cv2.imwrite(curr_name+'.png', images_aug[j,:])
                        self.image_path_list[part].append(curr_name+'.png')
                        self.label_list[part].append(curr_label)
                        curr_im += 1

            self.logger.info('Done image annotation %s', anno_name)

        return (im, anno)

    def end(self):

        EXEC_ID=str(int(time.time()))  # id used for saving files
        with open(OUTPUT_DIR+'/train/image_path_list_'+EXEC_ID+'.txt', 'w') as f:
            for p in self.image_path_list[0]:
                f.write(p+'\n')

        with open(OUTPUT_DIR+'/train/label_list_'+EXEC_ID+'.txt', 'w') as f:
            for l in self.label_list[0]:
                f.write(l+'\n')

        with open(OUTPUT_DIR+'/val/image_path_list_'+EXEC_ID+'.txt', 'w') as f:
            for p in self.image_path_list[1]:
                f.write(p+'\n')

        with open(OUTPUT_DIR+'/val/label_list_'+EXEC_ID+'.txt', 'w') as f:
            for l in self.label_list[1]:
                f.write(l+'\n')

        self.logger.info('Execution id: %s', EXEC_ID)

    def get_parameters(self):
        return {'parameters': []}

    def get_version(self):
        return self._VERSION

    def safe_scale_bbox(self, l, t, r, b, scale, shape):
        dw = (r-l)*(scale-1)/6.0
        dh = (b-t)*(scale-1)/2.0

        l = int(l-dw)
        t = int(t-dh)
        r = int(r+dw)
        b = int(b+dh)

        if l<0: l=0
        if t<0: t=0
        if b>=shape[0]: b=shape[0]-1
        if r>=shape[1]: r=shape[1]-1

        return (l,t,r,b)
<|endoftext|>


<commit_before>from __future__ import division, print_function

from functools import partial
from collections import defaultdict
import itertools

import numpy as np

from .munkres import linear_assignment

# TODO: Blanc and standard clustering metrics (e.g. http://scikit-learn.org/stable/modules/clustering.html)
# TODO: cite originating papers
# XXX: perhaps use set (or list) of sets rather than dict of sets

def mapping_to_sets(mapping):
    """
    >>> sets = mapping_to_sets({'a': 1, 'b': 2, 'c': 1}).items()
    >>> sorted((k, sorted(v)) for k, v in sets)
    [(1, ['a', 'c']), (2, ['b'])]
    """
    s = defaultdict(set)
    for m, k in mapping.items():
        s[k].add(m)
    return dict(s)

def sets_to_mapping(s):
    """
    >>> sorted(sets_to_mapping({1: {'a', 'c'}, 2: {'b'}}).items())
    [('a', 1), ('b', 2), ('c', 1)]
    """
    return {m: k for k, ms in s.items() for m in ms}

def _f1(a, b):
    return 2 * a * b / (a + b)

def dice(a, b):
    """

    "Entity-based" measure in CoNLL; #4 in CEAF paper
    """
    if a and b:
        return len(a & b) / (len(a) + len(b))
    return 0.

def overlap(a, b):
    """Intersection of sets

    "Mention-based" measure in CoNLL; #3 in CEAF paper
    """
    return len(a & b)

def ceaf(true, pred, similarity=dice):
    """

    >>> true = {'A': {1,2,3,4,5}, 'B': {6,7}, 'C': {8, 9, 10, 11, 12}}
    >>> pred_a = {'A': {1,2,3,4,5}, 'B': {6,7, 8, 9, 10, 11, 12}}
    >>> pred_b = {'A': {1,2,3,4,5,8, 9, 10, 11, 12}, 'B': {6,7}}
    >>> pred_c = {'A': {1,2,3,4,5, 6,7, 8, 9, 10, 11, 12}}
    >>> pred_d = {i: {i,} for i in range(1, 13)}
    >>> mention_ceaf(true, pred_a)[-1]  # doctest: +ELLIPSIS
    0.83...
    >>> entity_ceaf(true, pred_a)[-1]  # doctest: +ELLIPSIS
    0.73...
    >>> mention_ceaf(true, pred_b)[-1]  # doctest: +ELLIPSIS
    0.58...
    >>> entity_ceaf(true, pred_b)[-1]  # doctest: +ELLIPSIS
    0.66...
    >>> mention_ceaf(true, pred_c)  # doctest: +ELLIPSIS
    (0.416..., 0.416..., 0.416...)
    >>> entity_ceaf(true, pred_c)  # doctest: +ELLIPSIS
    (0.588..., 0.196..., 0.294...)
    >>> mention_ceaf(true, pred_d)  # doctest: +ELLIPSIS
    (0.25, 0.25, 0.25)
    >>> entity_ceaf(true, pred_d)  # doctest: +ELLIPSIS
    (0.111..., 0.444..., 0.177...)
    """
    X = np.empty((len(true), len(pred)))
    pred = list(pred.values())
    for R, Xrow in zip(true.values(), X):
        Xrow[:] = [similarity(R, S) for S in pred]
    indices = linear_assignment(-X)

    numerator = sum(X[indices[:, 0], indices[:, 1]])
    true_denom = sum(similarity(R, R) for R in true.values())
    pred_denom = sum(similarity(S, S) for S in pred)
    p = numerator / pred_denom
    r = numerator / true_denom
    return p, r, _f1(p, r)

entity_ceaf = partial(ceaf, similarity=dice)
mention_ceaf = partial(ceaf, similarity=overlap)

def _b_cubed(A, B, A_mapping, B_mapping, EMPTY=frozenset([])):
    res = 0.
    for m, k in A_mapping.items():
        A_cluster = A.get(k, EMPTY)
        res += len(A_cluster & B.get(B_mapping.get(m), EMPTY)) / len(A_cluster)
    res /= len(A_mapping)
    return res

def b_cubed(true, pred):
    """

    TODO: tests
    """
    true_mapping = sets_to_mapping(true)
    pred_mapping = sets_to_mapping(pred)
    p = _b_cubed(pred, true, pred_mapping, true_mapping)
    r = _b_cubed(true, pred, true_mapping, pred_mapping)
    return p, r, _f1(p, r)

def pairwise_f1(true, pred):
    """Measure the proportion of correctly identified pairwise coindexations

    TODO: tests
    """
    pred_mapping = sets_to_mapping(pred)
    correct = 0
    for cluster in true.values():
        for m1, m2 in itertools.combinations(cluster, 2):
            if pred_mapping.get(m1) == pred_mapping.get(m2):
                correct += 1
    p = correct / sum(len(cluster) - 1 for cluster in pred.values())
    r = correct / sum(len(cluster) - 1 for cluster in true.values())
    return p, r, _f1(p, r)

def _vilain(A, B_mapping):
    numerator = 0
    denominator = 0
    for cluster in A.values():
        corresponding = set()
        n_unaligned = 0
        for m in cluster:
            if m not in B_mapping:
                n_unaligned += 1
            else:
                corresponding.add(B_mapping[m])
        numerator += len(cluster) - n_unaligned - len(corresponding)
        denominator += len(cluster) - 1
    return numerator / denominator

def muc(true, pred):
    """The MUC evaluation metric defined in Vilain et al. (1995)

    This calculates recall error for each true cluster C as the number of
    response clusters that would need to be merged in order to produce a
    superset of C.

    Examples from Vilain et al. (1995):
    >>> muc({1: {'A', 'B', 'C', 'D'}},
    ...     {1: {'A', 'B'}, 2: {'C', 'D'}})  # doctest: +ELLIPSIS
    (1.0, 0.66..., 0.8)
    >>> muc({1: {'A', 'B'}, 2: {'C', 'D'}},
    ...     {1: {'A', 'B', 'C', 'D'}})  # doctest: +ELLIPSIS
    (0.66..., 1.0, 0.8)
    >>> muc({1: {'A', 'B', 'C'}}, {1: {'A', 'C'}})  # doctest: +ELLIPSIS
    (1.0, 0.5, 0.66...)
    >>> muc({1: {'B', 'C', 'D', 'E', 'G', 'H', 'J'}},
    ...     {1: {'A', 'B', 'C'}, 2: {'D', 'E', 'F'}, 3: {'G', 'H', 'I'}})
    ... # doctest: +ELLIPSIS
    (0.5, 0.5, 0.5)
    >>> muc({1: {'A', 'B', 'C'}, 2: {'D', 'E', 'F', 'G'}},
    ...     {1: {'A', 'B'}, 2: {'C', 'D'}, 3: {'F', 'G', 'H'}})
    ... # doctest: +ELLIPSIS
    (0.5, 0.4, 0.44...)
    """
    p = _vilain(pred, sets_to_mapping(true))
    r = _vilain(true, sets_to_mapping(pred))
    return p, r, _f1(p, r)

def read_conll_coref(f):
    res = defaultdict(set)
    # TODO: handle annotations over document boundary
    start = None
    i = 0
    for l in f:
        if l.startswith('#') or not l.strip():
            continue
        i += 1
        doc_id, tag = l.strip().split(' ', 1)
        if tag == '-':
            continue
        if tag.endswith(')'):
            if start is None:
                assert tag.startswith('(')
            else:
                assert not tag.startswith('(')
            cid = tag.lstrip('(').rstrip(')')
            res[cid].add((doc_id, start, i))
            start = None
        elif tag.startswith('('):
            start = i
    return dict(res)

if __name__ == '__main__':
    import argparse
    ap = argparse.ArgumentParser(description='CoNLL2011-2 coreference evaluator')
    ap.add_argument('key_file', type=argparse.FileType('r'))
    ap.add_argument('response_file', type=argparse.FileType('r'))
    args = ap.parse_args()
    METRICS = {
        'bcubed': b_cubed,
        'ceafe': entity_ceaf,
        'ceafm': mention_ceaf,
        'muc': muc,
        'pairs': pairwise_f1,
    }
    key = read_conll_coref(args.key_file)
    response = read_conll_coref(args.response_file)
    print('Metric', 'P', 'R', 'F1', sep='\t')
    for name, fn in sorted(METRICS.items()):
        print(name, *('{:0.2f}'.format(100 * x) for x in fn(key, response)), sep='\t')
<commit_msg>Fix conll reading, pairwise denominator<commit_after>from __future__ import division, print_function

from functools import partial
from collections import defaultdict
import itertools
import re

import numpy as np

from .munkres import linear_assignment

# TODO: Blanc and standard clustering metrics (e.g. http://scikit-learn.org/stable/modules/clustering.html)
# TODO: cite originating papers
# XXX: perhaps use set (or list) of sets rather than dict of sets

def mapping_to_sets(mapping):
    """
    >>> sets = mapping_to_sets({'a': 1, 'b': 2, 'c': 1}).items()
    >>> sorted((k, sorted(v)) for k, v in sets)
    [(1, ['a', 'c']), (2, ['b'])]
    """
    s = defaultdict(set)
    for m, k in mapping.items():
        s[k].add(m)
    return dict(s)

def sets_to_mapping(s):
    """
    >>> sorted(sets_to_mapping({1: {'a', 'c'}, 2: {'b'}}).items())
    [('a', 1), ('b', 2), ('c', 1)]
    """
    return {m: k for k, ms in s.items() for m in ms}

def _f1(a, b):
    return 2 * a * b / (a + b)

def dice(a, b):
    """

    "Entity-based" measure in CoNLL; #4 in CEAF paper
    """
    if a and b:
        return len(a & b) / (len(a) + len(b))
    return 0.

def overlap(a, b):
    """Intersection of sets

    "Mention-based" measure in CoNLL; #3 in CEAF paper
    """
    return len(a & b)

def ceaf(true, pred, similarity=dice):
    """

    >>> true = {'A': {1,2,3,4,5}, 'B': {6,7}, 'C': {8, 9, 10, 11, 12}}
    >>> pred_a = {'A': {1,2,3,4,5}, 'B': {6,7, 8, 9, 10, 11, 12}}
    >>> pred_b = {'A': {1,2,3,4,5,8, 9, 10, 11, 12}, 'B': {6,7}}
    >>> pred_c = {'A': {1,2,3,4,5, 6,7, 8, 9, 10, 11, 12}}
    >>> pred_d = {i: {i,} for i in range(1, 13)}
    >>> mention_ceaf(true, pred_a)[-1]  # doctest: +ELLIPSIS
    0.83...
    >>> entity_ceaf(true, pred_a)[-1]  # doctest: +ELLIPSIS
    0.73...
    >>> mention_ceaf(true, pred_b)[-1]  # doctest: +ELLIPSIS
    0.58...
    >>> entity_ceaf(true, pred_b)[-1]  # doctest: +ELLIPSIS
    0.66...
    >>> mention_ceaf(true, pred_c)  # doctest: +ELLIPSIS
    (0.416..., 0.416..., 0.416...)
    >>> entity_ceaf(true, pred_c)  # doctest: +ELLIPSIS
    (0.588..., 0.196..., 0.294...)
    >>> mention_ceaf(true, pred_d)  # doctest: +ELLIPSIS
    (0.25, 0.25, 0.25)
    >>> entity_ceaf(true, pred_d)  # doctest: +ELLIPSIS
    (0.111..., 0.444..., 0.177...)
    """
    X = np.empty((len(true), len(pred)))
    pred = list(pred.values())
    for R, Xrow in zip(true.values(), X):
        Xrow[:] = [similarity(R, S) for S in pred]
    indices = linear_assignment(-X)

    numerator = sum(X[indices[:, 0], indices[:, 1]])
    true_denom = sum(similarity(R, R) for R in true.values())
    pred_denom = sum(similarity(S, S) for S in pred)
    p = numerator / pred_denom
    r = numerator / true_denom
    return p, r, _f1(p, r)

entity_ceaf = partial(ceaf, similarity=dice)
mention_ceaf = partial(ceaf, similarity=overlap)

def _b_cubed(A, B, A_mapping, B_mapping, EMPTY=frozenset([])):
    res = 0.
    for m, k in A_mapping.items():
        A_cluster = A.get(k, EMPTY)
        res += len(A_cluster & B.get(B_mapping.get(m), EMPTY)) / len(A_cluster)
    res /= len(A_mapping)
    return res

def b_cubed(true, pred):
    """

    TODO: tests
    """
    true_mapping = sets_to_mapping(true)
    pred_mapping = sets_to_mapping(pred)
    p = _b_cubed(pred, true, pred_mapping, true_mapping)
    r = _b_cubed(true, pred, true_mapping, pred_mapping)
    return p, r, _f1(p, r)

def pairwise_f1(true, pred):
    """Measure the proportion of correctly identified pairwise coindexations

    TODO: tests
    """
    pred_mapping = sets_to_mapping(pred)
    correct = 0
    for cluster in true.values():
        for m1, m2 in itertools.combinations(cluster, 2):
            if pred_mapping.get(m1) == pred_mapping.get(m2):
                correct += 1
    p = correct / sum(len(cluster) * (len(cluster) - 1) for cluster in pred.values()) * 2
    r = correct / sum(len(cluster) * (len(cluster) - 1) for cluster in true.values()) * 2
    return p, r, _f1(p, r)

def _vilain(A, B_mapping):
    numerator = 0
    denominator = 0
    for cluster in A.values():
        corresponding = set()
        n_unaligned = 0
        for m in cluster:
            if m not in B_mapping:
                n_unaligned += 1
            else:
                corresponding.add(B_mapping[m])
        numerator += len(cluster) - n_unaligned - len(corresponding)
        denominator += len(cluster) - 1
    return numerator / denominator

def muc(true, pred):
    """The MUC evaluation metric defined in Vilain et al. (1995)

    This calculates recall error for each true cluster C as the number of
    response clusters that would need to be merged in order to produce a
    superset of C.

    Examples from Vilain et al. (1995):
    >>> muc({1: {'A', 'B', 'C', 'D'}},
    ...     {1: {'A', 'B'}, 2: {'C', 'D'}})  # doctest: +ELLIPSIS
    (1.0, 0.66..., 0.8)
    >>> muc({1: {'A', 'B'}, 2: {'C', 'D'}},
    ...     {1: {'A', 'B', 'C', 'D'}})  # doctest: +ELLIPSIS
    (0.66..., 1.0, 0.8)
    >>> muc({1: {'A', 'B', 'C'}}, {1: {'A', 'C'}})  # doctest: +ELLIPSIS
    (1.0, 0.5, 0.66...)
    >>> muc({1: {'B', 'C', 'D', 'E', 'G', 'H', 'J'}},
    ...     {1: {'A', 'B', 'C'}, 2: {'D', 'E', 'F'}, 3: {'G', 'H', 'I'}})
    ... # doctest: +ELLIPSIS
    (0.5, 0.5, 0.5)
    >>> muc({1: {'A', 'B', 'C'}, 2: {'D', 'E', 'F', 'G'}},
    ...     {1: {'A', 'B'}, 2: {'C', 'D'}, 3: {'F', 'G', 'H'}})
    ... # doctest: +ELLIPSIS
    (0.5, 0.4, 0.44...)
    """
    p = _vilain(pred, sets_to_mapping(true))
    r = _vilain(true, sets_to_mapping(pred))
    return p, r, _f1(p, r)

def read_conll_coref(f):
    res = defaultdict(set)
    # TODO: handle annotations over document boundary
    i = 0
    opened = {}
    for l in f:
        if l.startswith('#'):
            continue
        l = l.split()
        if not l:
            assert not opened
            continue

        i += 1
        tag = l[-1]

        for match in re.finditer(r'\(?[0-9]+\)?', tag):
            match = match.group()
            cid = match.strip('()')
            if match.startswith('('):
                assert cid not in opened
                opened[cid] = i
            if match.endswith(')'):
                res[cid].add((opened.pop(cid), i))
    return dict(res)

if __name__ == '__main__':
    import argparse
    ap = argparse.ArgumentParser(description='CoNLL2011-2 coreference evaluator')
    ap.add_argument('key_file', type=argparse.FileType('r'))
    ap.add_argument('response_file', type=argparse.FileType('r'))
    args = ap.parse_args()
    METRICS = {
        'bcubed': b_cubed,
        'ceafe': entity_ceaf,
        'ceafm': mention_ceaf,
        'muc': muc,
        'pairs': pairwise_f1,
    }
    key = read_conll_coref(args.key_file)
    response = read_conll_coref(args.response_file)
    print('Metric', 'P', 'R', 'F1', sep='\t')
    for name, fn in sorted(METRICS.items()):
        print(name, *('{:0.2f}'.format(100 * x) for x in fn(key, response)), sep='\t')
<|endoftext|>


<commit_before>from __future__ import division, absolute_import, print_function,\
 unicode_literals

import ctypes
import platform
import sys

if sys.platform in ('win32', 'cygwin'):
    _functype = ctypes.WINFUNCTYPE
    _lib = ctypes.windll.nanomsg
else:
    _functype = ctypes.CFUNCTYPE
    _lib = ctypes.cdll.LoadLibrary('libnanomsg.so')

def _c_func_wrapper_factory(cdecl_text):
    def move_pointer_and_strip(type_def, name):
        if '*' in name:
            type_def += ' ' + name[:name.rindex('*')+1]
            name = name.rsplit('*', 1)[1]
        return type_def.strip(), name.strip()

    def type_lookup(type_def):
        types = {
            'void': None,
            'char *': ctypes.c_char_p,
            'int': ctypes.c_int,
            'int *': ctypes.POINTER(ctypes.c_int),
            'void *': ctypes.c_void_p,
            'size_t':  ctypes.c_size_t,
            'size_t *':  ctypes.POINTER(ctypes.c_size_t),
            'struct nn_msghdr *': ctypes.c_void_p,
        }
        type_def_without_const = type_def.replace('const ','')
        if type_def_without_const in types:
            return types[type_def_without_const]
        elif (type_def_without_const.endswith('*') and
                type_def_without_const[:-1] in types):
            return ctypes.POINTER(types[type_def_without_const[:-1]])
        else:
            raise KeyError(type_def)

        return types[type_def.replace('const ','')]

    a, b = [i.strip() for i in cdecl_text.split('(',1)]
    params, _ = b.rsplit(')',1)
    rtn_type, name = move_pointer_and_strip(*a.rsplit(' ', 1))
    param_spec = []
    for param in params.split(','):
        if param != 'void':
            param_spec.append(move_pointer_and_strip(*param.rsplit(' ', 1)))
    func = _functype(type_lookup(rtn_type),
                     *[type_lookup(type_def) for type_def, _ in param_spec])(
                        (name, _lib),
                        tuple((2 if '**' in type_def else 1, name)
                              for type_def, name in param_spec)
                    )
    func.__name__ = name
    return func

_C_HEADER = """
NN_EXPORT int nn_errno (void);
NN_EXPORT const char *nn_strerror (int errnum);
NN_EXPORT const char *nn_symbol (int i, int *value);
NN_EXPORT void nn_term (void);
NN_EXPORT void *nn_allocmsg (size_t size, int type);
NN_EXPORT int nn_freemsg (void *msg);
NN_EXPORT int nn_socket (int domain, int protocol);
NN_EXPORT int nn_close (int s);
NN_EXPORT int nn_setsockopt (int s, int level, int option, const void \
*optval, size_t optvallen);
NN_EXPORT int nn_getsockopt (int s, int level, int option, void *optval, \
size_t *optvallen);
NN_EXPORT int nn_bind (int s, const char *addr);
NN_EXPORT int nn_connect (int s, const char *addr);
NN_EXPORT int nn_shutdown (int s, int how);
NN_EXPORT int nn_send (int s, const void *buf, size_t len, int flags);
NN_EXPORT int nn_recv (int s, void *buf, size_t len, int flags);
NN_EXPORT int nn_sendmsg (int s, const struct nn_msghdr *msghdr, int flags);
NN_EXPORT int nn_recvmsg (int s, struct nn_msghdr *msghdr, int flags);
NN_EXPORT int nn_device (int s1, int s2);\
""".replace('NN_EXPORT', '')

for cdecl_text in _C_HEADER.splitlines():
    if cdecl_text.strip():
        func = _c_func_wrapper_factory(cdecl_text)
        globals()['_' + func.__name__] = func

def nn_symbols():
    "query the names and values of nanomsg symbols"
    value = ctypes.c_int()
    name_value_pairs = []
    i = 0
    while True:
        name = _nn_symbol(i, ctypes.byref(value))
        if name is None:
            break
        i += 1
        name_value_pairs.append((name.decode('ascii'), value.value))
    return name_value_pairs

nn_errno = _nn_errno
nn_errno.__doc__ = "retrieve the current errno"

nn_strerror = _nn_strerror
nn_strerror.__doc__ = "convert an error number into human-readable string"

nn_socket = _nn_socket
nn_socket.__doc__ = "create an SP socket"

nn_close = _nn_close
nn_close.__doc__ = "close an SP socket"

nn_bind = _nn_bind
nn_bind.__doc__ = "add a local endpoint to the socket"

nn_connect = _nn_connect
nn_connect.__doc__ = "add a remote endpoint to the socket"

nn_shutdown = _nn_shutdown
nn_shutdown.__doc__ = "remove an endpoint from a socket"

def create_writable_buffer(size):
    """Returns a writable buffer.

    This is the ctypes implementation.
    """
    return (ctypes.c_ubyte*size)()

def nn_setsockopt(socket, level, option, value):
    """set a socket option

    socket - socket number
    level - option level
    option - option
    value - a readable byte buffer (not a Unicode string) containing the value
    returns - 0 on success or < 0 on error

    """
    try:
        return _nn_setsockopt(socket, level, option, ctypes.addressof(value),
                              len(value))
    except (TypeError, AttributeError):
        buf_value = ctypes.create_string_buffer(value)
        return _nn_setsockopt(socket, level, option,
                              ctypes.addressof(buf_value), len(value))

def nn_getsockopt(socket, level, option, value):
    """retrieve a socket option

    socket - socket number
    level - option level
    option - option
    value - a writable byte buffer (e.g. a bytearray) which the option value
    will be copied to
    returns - number of bytes copied or on error nunber < 0

    """
    if memoryview(value).readonly:
        raise TypeError('Writable buffer is required')
    size_t_size = ctypes.c_size_t(len(value))
    rtn = _nn_getsockopt(socket, level, option, ctypes.addressof(value),
                         ctypes.byref(size_t_size))
    return (rtn, size_t_size.value)

def nn_send(socket, msg, flags):
    "send a message"
    try:
        return _nn_send(socket, ctypes.addressof(msg), len(buffer(msg)), flags)
    except (TypeError, AttributeError):
        buf_msg = ctypes.create_string_buffer(msg)
        return _nn_send(socket, ctypes.addressof(buf_msg), len(msg), flags)

def _create_message(address, length):
    class Message(ctypes.Union):
        _fields_ = [('_buf', ctypes.c_ubyte*length)]
        _len = length
        _address = address

        def __repr__(self):
            return '<_nanomsg_cpy.Message size %d, address 0x%x >' % (
                self._len,
                self._address
            )

        def __str__(self):
            return bytes(buffer(self))

        def __del__(self):
            _nn_freemsg(self._address)
            self._len = 0
            self._address = 0

    return Message.from_address(address)

def nn_allocmsg(size, type):
    "allocate a message"
    pointer = _nn_allocmsg(size, type)
    if pointer is None:
        return None
    return _create_message(pointer, size)

def nn_recv(socket, *args):
    "receive a message"
    if len(args) == 1:
        flags, = args
        pointer = ctypes.c_void_p()
        rtn = _nn_recv(socket, ctypes.byref(pointer), ctypes.c_size_t(-1),
                       flags)
        if rtn < 0:
            return rtn, None
        else:
            return rtn, _create_message(pointer.value, rtn)
    elif len(args) == 2:
        msg_buf, flags = args
        mv_buf = memoryview(msg_buf)
        if mv_buf.readonly:
            raise TypeError('Writable buffer is required')
        rtn = _nn_recv(socket, ctypes.addressof(msg_buf), len(mv_buf), flags)
        return rtn, msg_buf

nn_device = _nn_device
nn_device.__doc__ = "start a device"

nn_term = _nn_term
nn_term.__doc__ = "notify all sockets about process termination"

try:
    if sys.platform in ('win32', 'cygwin'):
        _nclib = ctypes.windll.nanoconfig
    else:
        _nclib = ctypes.cdll.LoadLibrary('libnanoconfig.so')
except OSError:
    pass # No nanoconfig, sorry
else:
    # int nc_configure (int s, const char *addr)
    nc_configure = _functype(ctypes.c_int, ctypes.c_int, ctypes.c_char_p)(
        ('nc_configure', _nclib), ((1, 's'), (1, 'addr')))
    nc_configure.__doc__ = "configure socket using nanoconfig"

    # void nc_close(int s);
    nc_close = _functype(None, ctypes.c_int)(
        ('nc_close', _nclib), ((1, 's'),))
    nc_close.__doc__ = "close an SP socket configured with nn_configure"

    # void nc_term();
    nc_term = _functype(None)(
        ('nc_term', _nclib), ())
    nc_term.__doc__ = "shutdown nanoconfig worker thread"
<commit_msg>Added osx support to ctypes wrapper<commit_after>from __future__ import division, absolute_import, print_function,\
 unicode_literals

import ctypes
import platform
import sys

if sys.platform in ('win32', 'cygwin'):
    _functype = ctypes.WINFUNCTYPE
    _lib = ctypes.windll.nanomsg
elif sys.platform in ('darwin'):
    _functype = ctypes.CFUNCTYPE
    _lib = ctypes.cdll.LoadLibrary('libnanomsg.dylib')
else:
    _functype = ctypes.CFUNCTYPE
    _lib = ctypes.cdll.LoadLibrary('libnanomsg.so')

def _c_func_wrapper_factory(cdecl_text):
    def move_pointer_and_strip(type_def, name):
        if '*' in name:
            type_def += ' ' + name[:name.rindex('*')+1]
            name = name.rsplit('*', 1)[1]
        return type_def.strip(), name.strip()

    def type_lookup(type_def):
        types = {
            'void': None,
            'char *': ctypes.c_char_p,
            'int': ctypes.c_int,
            'int *': ctypes.POINTER(ctypes.c_int),
            'void *': ctypes.c_void_p,
            'size_t':  ctypes.c_size_t,
            'size_t *':  ctypes.POINTER(ctypes.c_size_t),
            'struct nn_msghdr *': ctypes.c_void_p,
        }
        type_def_without_const = type_def.replace('const ','')
        if type_def_without_const in types:
            return types[type_def_without_const]
        elif (type_def_without_const.endswith('*') and
                type_def_without_const[:-1] in types):
            return ctypes.POINTER(types[type_def_without_const[:-1]])
        else:
            raise KeyError(type_def)

        return types[type_def.replace('const ','')]

    a, b = [i.strip() for i in cdecl_text.split('(',1)]
    params, _ = b.rsplit(')',1)
    rtn_type, name = move_pointer_and_strip(*a.rsplit(' ', 1))
    param_spec = []
    for param in params.split(','):
        if param != 'void':
            param_spec.append(move_pointer_and_strip(*param.rsplit(' ', 1)))
    func = _functype(type_lookup(rtn_type),
                     *[type_lookup(type_def) for type_def, _ in param_spec])(
                        (name, _lib),
                        tuple((2 if '**' in type_def else 1, name)
                              for type_def, name in param_spec)
                    )
    func.__name__ = name
    return func

_C_HEADER = """
NN_EXPORT int nn_errno (void);
NN_EXPORT const char *nn_strerror (int errnum);
NN_EXPORT const char *nn_symbol (int i, int *value);
NN_EXPORT void nn_term (void);
NN_EXPORT void *nn_allocmsg (size_t size, int type);
NN_EXPORT int nn_freemsg (void *msg);
NN_EXPORT int nn_socket (int domain, int protocol);
NN_EXPORT int nn_close (int s);
NN_EXPORT int nn_setsockopt (int s, int level, int option, const void \
*optval, size_t optvallen);
NN_EXPORT int nn_getsockopt (int s, int level, int option, void *optval, \
size_t *optvallen);
NN_EXPORT int nn_bind (int s, const char *addr);
NN_EXPORT int nn_connect (int s, const char *addr);
NN_EXPORT int nn_shutdown (int s, int how);
NN_EXPORT int nn_send (int s, const void *buf, size_t len, int flags);
NN_EXPORT int nn_recv (int s, void *buf, size_t len, int flags);
NN_EXPORT int nn_sendmsg (int s, const struct nn_msghdr *msghdr, int flags);
NN_EXPORT int nn_recvmsg (int s, struct nn_msghdr *msghdr, int flags);
NN_EXPORT int nn_device (int s1, int s2);\
""".replace('NN_EXPORT', '')

for cdecl_text in _C_HEADER.splitlines():
    if cdecl_text.strip():
        func = _c_func_wrapper_factory(cdecl_text)
        globals()['_' + func.__name__] = func

def nn_symbols():
    "query the names and values of nanomsg symbols"
    value = ctypes.c_int()
    name_value_pairs = []
    i = 0
    while True:
        name = _nn_symbol(i, ctypes.byref(value))
        if name is None:
            break
        i += 1
        name_value_pairs.append((name.decode('ascii'), value.value))
    return name_value_pairs

nn_errno = _nn_errno
nn_errno.__doc__ = "retrieve the current errno"

nn_strerror = _nn_strerror
nn_strerror.__doc__ = "convert an error number into human-readable string"

nn_socket = _nn_socket
nn_socket.__doc__ = "create an SP socket"

nn_close = _nn_close
nn_close.__doc__ = "close an SP socket"

nn_bind = _nn_bind
nn_bind.__doc__ = "add a local endpoint to the socket"

nn_connect = _nn_connect
nn_connect.__doc__ = "add a remote endpoint to the socket"

nn_shutdown = _nn_shutdown
nn_shutdown.__doc__ = "remove an endpoint from a socket"

def create_writable_buffer(size):
    """Returns a writable buffer.

    This is the ctypes implementation.
    """
    return (ctypes.c_ubyte*size)()

def nn_setsockopt(socket, level, option, value):
    """set a socket option

    socket - socket number
    level - option level
    option - option
    value - a readable byte buffer (not a Unicode string) containing the value
    returns - 0 on success or < 0 on error

    """
    try:
        return _nn_setsockopt(socket, level, option, ctypes.addressof(value),
                              len(value))
    except (TypeError, AttributeError):
        buf_value = ctypes.create_string_buffer(value)
        return _nn_setsockopt(socket, level, option,
                              ctypes.addressof(buf_value), len(value))

def nn_getsockopt(socket, level, option, value):
    """retrieve a socket option

    socket - socket number
    level - option level
    option - option
    value - a writable byte buffer (e.g. a bytearray) which the option value
    will be copied to
    returns - number of bytes copied or on error nunber < 0

    """
    if memoryview(value).readonly:
        raise TypeError('Writable buffer is required')
    size_t_size = ctypes.c_size_t(len(value))
    rtn = _nn_getsockopt(socket, level, option, ctypes.addressof(value),
                         ctypes.byref(size_t_size))
    return (rtn, size_t_size.value)

def nn_send(socket, msg, flags):
    "send a message"
    try:
        return _nn_send(socket, ctypes.addressof(msg), len(buffer(msg)), flags)
    except (TypeError, AttributeError):
        buf_msg = ctypes.create_string_buffer(msg)
        return _nn_send(socket, ctypes.addressof(buf_msg), len(msg), flags)

def _create_message(address, length):
    class Message(ctypes.Union):
        _fields_ = [('_buf', ctypes.c_ubyte*length)]
        _len = length
        _address = address

        def __repr__(self):
            return '<_nanomsg_cpy.Message size %d, address 0x%x >' % (
                self._len,
                self._address
            )

        def __str__(self):
            return bytes(buffer(self))

        def __del__(self):
            _nn_freemsg(self._address)
            self._len = 0
            self._address = 0

    return Message.from_address(address)

def nn_allocmsg(size, type):
    "allocate a message"
    pointer = _nn_allocmsg(size, type)
    if pointer is None:
        return None
    return _create_message(pointer, size)

def nn_recv(socket, *args):
    "receive a message"
    if len(args) == 1:
        flags, = args
        pointer = ctypes.c_void_p()
        rtn = _nn_recv(socket, ctypes.byref(pointer), ctypes.c_size_t(-1),
                       flags)
        if rtn < 0:
            return rtn, None
        else:
            return rtn, _create_message(pointer.value, rtn)
    elif len(args) == 2:
        msg_buf, flags = args
        mv_buf = memoryview(msg_buf)
        if mv_buf.readonly:
            raise TypeError('Writable buffer is required')
        rtn = _nn_recv(socket, ctypes.addressof(msg_buf), len(mv_buf), flags)
        return rtn, msg_buf

nn_device = _nn_device
nn_device.__doc__ = "start a device"

nn_term = _nn_term
nn_term.__doc__ = "notify all sockets about process termination"

try:
    if sys.platform in ('win32', 'cygwin'):
        _nclib = ctypes.windll.nanoconfig
    else:
        _nclib = ctypes.cdll.LoadLibrary('libnanoconfig.so')
except OSError:
    pass # No nanoconfig, sorry
else:
    # int nc_configure (int s, const char *addr)
    nc_configure = _functype(ctypes.c_int, ctypes.c_int, ctypes.c_char_p)(
        ('nc_configure', _nclib), ((1, 's'), (1, 'addr')))
    nc_configure.__doc__ = "configure socket using nanoconfig"

    # void nc_close(int s);
    nc_close = _functype(None, ctypes.c_int)(
        ('nc_close', _nclib), ((1, 's'),))
    nc_close.__doc__ = "close an SP socket configured with nn_configure"

    # void nc_term();
    nc_term = _functype(None)(
        ('nc_term', _nclib), ())
    nc_term.__doc__ = "shutdown nanoconfig worker thread"
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import elasticsearch
import geopy.geocoders
import geopy.exc
import transaction
import logging

import pyjobsweb.commands
import pyjobsweb.model
import pyjobsweb.lib

class PopulateESCommand(pyjobsweb.commands.AppContextCommand):
    @staticmethod
    def __compute_pending_insertions():
        pending_insertions = pyjobsweb.model.DBSession\
            .query(pyjobsweb.model.data.JobOfferSQLAlchemy)\
            .filter_by(
                already_in_elasticsearch=False
            )
        return pending_insertions

    @staticmethod
    def __mark_task_as_handled(job_offer):
        transaction.begin()
        pyjobsweb.model.DBSession\
            .query(pyjobsweb.model.data.JobOfferSQLAlchemy)\
            .filter(pyjobsweb.model.data.JobOfferSQLAlchemy.id == job_offer.id)\
            .update({'already_in_elasticsearch': True})
        transaction.commit()

    def __handle_insertion_task(self, job_offer):
        es_job_offer = job_offer.to_elasticsearch_job_offer()

        # Perform the insertion in Elasticsearch
        try:
            # Compute lat, lng of the job offer
            geolocator = geopy.geocoders.Nominatim(timeout=5)
            location = geolocator.geocode(job_offer.address, True)

            if not location:
                raise PopulateESCommand.FailedGeoloc

            es_job_offer.geolocation = [location.latitude, location.longitude]
            es_job_offer.geolocation_error = False
        except (geopy.exc.GeocoderQueryError, PopulateESCommand.FailedGeoloc):
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Couldn't compute geolocation of the following address : ",
                    u''.join(job_offer.address).encode('utf-8').strip()
                )
            )
            es_job_offer.geolocation = [0, 0]
            es_job_offer.geolocation_error = True
        except (geopy.exc.GeocoderQuotaExceeded,
                geopy.exc.GeocoderAuthenticationFailure,
                geopy.exc.GeocoderUnavailable,
                geopy.exc.GeocoderNotFound) as e:
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Error during geolocation : ",
                    e.message,
                    ". Aborting Elasticsearch insertions now..."
                )
            )
            raise PopulateESCommand.AbortException
        except geopy.exc.GeopyError as e:
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Error during geolocation : ",
                    e.message
                )
            )

        try:
            # Perform the insertion
            es_job_offer.save()
        except elasticsearch.exceptions.RequestError as e:
            logging.getLogger(__name__).log(
                logging.ERROR,
                '{}{}{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    ", error during the Elasticsearch insertion : ",
                    e,
                    ".\nInfo : ",
                    e.info
                )
            )
            return

        # Mark the task as handled so we don't retreat it next time
        self.__mark_task_as_handled(job_offer)

    def take_action(self, parsed_args):
        super(PopulateESCommand, self).take_action(parsed_args)

        job_offers = self.__compute_pending_insertions()

        for j in job_offers:
            try:
                self.__handle_insertion_task(j)
            except PopulateESCommand.AbortException:
                return

    class FailedGeoloc(Exception):
        pass

    class AbortException(Exception):
        pass
<commit_msg>Bugfix geolocation storage<commit_after># -*- coding: utf-8 -*-
import elasticsearch
import geopy.geocoders
import geopy.exc
import transaction
import logging

import pyjobsweb.commands
import pyjobsweb.model
import pyjobsweb.lib

class PopulateESCommand(pyjobsweb.commands.AppContextCommand):
    @staticmethod
    def __compute_pending_insertions():
        pending_insertions = pyjobsweb.model.DBSession\
            .query(pyjobsweb.model.data.JobOfferSQLAlchemy)\
            .filter_by(
                already_in_elasticsearch=False
            )
        return pending_insertions

    @staticmethod
    def __mark_task_as_handled(job_offer):
        transaction.begin()
        pyjobsweb.model.DBSession\
            .query(pyjobsweb.model.data.JobOfferSQLAlchemy)\
            .filter(pyjobsweb.model.data.JobOfferSQLAlchemy.id == job_offer.id)\
            .update({'already_in_elasticsearch': True})
        transaction.commit()

    def __handle_insertion_task(self, job_offer):
        es_job_offer = job_offer.to_elasticsearch_job_offer()

        # Perform the insertion in Elasticsearch
        try:
            # Compute lat, lng of the job offer
            geolocator = geopy.geocoders.Nominatim(timeout=5)
            location = geolocator.geocode(job_offer.address, True)

            if not location:
                raise PopulateESCommand.FailedGeoloc

            es_job_offer.geolocation = [location.longitude, location.latitude]
            es_job_offer.geolocation_error = False
        except (geopy.exc.GeocoderQueryError, PopulateESCommand.FailedGeoloc):
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Couldn't compute geolocation of the following address : ",
                    u''.join(job_offer.address).encode('utf-8').strip()
                )
            )
            es_job_offer.geolocation = [0, 0]
            es_job_offer.geolocation_error = True
        except (geopy.exc.GeocoderQuotaExceeded,
                geopy.exc.GeocoderAuthenticationFailure,
                geopy.exc.GeocoderUnavailable,
                geopy.exc.GeocoderNotFound) as e:
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Error during geolocation : ",
                    e.message,
                    ". Aborting Elasticsearch insertions now..."
                )
            )
            raise PopulateESCommand.AbortException
        except geopy.exc.GeopyError as e:
            logging.getLogger(__name__).log(
                logging.WARNING,
                '{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    " Error during geolocation : ",
                    e.message
                )
            )

        try:
            # Perform the insertion
            es_job_offer.save()
        except elasticsearch.exceptions.RequestError as e:
            logging.getLogger(__name__).log(
                logging.ERROR,
                '{}{}{}{}{}{}'.format(
                    "Job offer id : ",
                    job_offer.id,
                    ", error during the Elasticsearch insertion : ",
                    e,
                    ".\nInfo : ",
                    e.info
                )
            )
            return

        # Mark the task as handled so we don't retreat it next time
        self.__mark_task_as_handled(job_offer)

    def take_action(self, parsed_args):
        super(PopulateESCommand, self).take_action(parsed_args)

        job_offers = self.__compute_pending_insertions()

        for j in job_offers:
            try:
                self.__handle_insertion_task(j)
            except PopulateESCommand.AbortException:
                return

    class FailedGeoloc(Exception):
        pass

    class AbortException(Exception):
        pass
<|endoftext|>


<commit_before>import random, hmac, hashlib
from django.db import models
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic
from django.core.urlresolvers import reverse
from utils import int_to_base32, base32_to_int, send_email

class EmailConfirmationManager(models.Manager):
    def confirm(self, request, object, after_confirm, after_unsubscribe):
        conf = EmailConfirmation(
            content_object=object,
            after_confirm=after_confirm,
            after_unsubscribe=after_unsubscribe,
        )
        conf.send_email(request)

    def path_for_unsubscribe(self, object):
        obj = EmailConfirmation.objects.get(content_object=object)
        return reverse('emailconfirmation-unsubscribe', kwargs={
            'id': int_to_base32(obj.id),
            'token': obj.make_token(random.randint(0,32767)),
        })

class EmailConfirmation(models.Model):
    confirmed = models.BooleanField(default=False)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')
    after_confirm = models.CharField(max_length=100)
    after_unsubscribe = models.CharField(max_length=100)

    objects = EmailConfirmationManager()

    def __unicode__(self):
        return 'Confirming of %s, %s' % (self.content_object.email, self.confirmed)

    def send_email(self, request):
        self.save()
        send_email(request, "Alert confirmation",
            'emailconfirmation/email.txt',
            {
                'object': self.content_object,
                'id': int_to_base32(self.id),
                'token': self.make_token(random.randint(0,32767)),
            }, self.content_object.email
        )

    @models.permalink
    def url_after_confirm(self):
        return (self.after_confirm, [ self.content_object.id ])

    @models.permalink
    def url_after_unsubscribe(self):
        return (self.after_unsubscribe, [ self.content_object.id ])

    def check_token(self, token):
        try:
            rand, hash = token.split("-")
        except:
            return False

        try:
            rand = base32_to_int(rand)
        except:
            return False

        if self.make_token(rand) != token:
            return False

        return True

    def make_token(self, rand):
        rand = int_to_base32(rand)
        hash = hmac.new(settings.SECRET_KEY, unicode(self.id) + rand, hashlib.sha1).hexdigest()[::2]
        return "%s-%s" % (rand, hash)
<commit_msg>Put function on model, not manager.<commit_after>import random, hmac, hashlib
from django.db import models
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic
from django.core.urlresolvers import reverse
from utils import int_to_base32, base32_to_int, send_email

class EmailConfirmationManager(models.Manager):
    def confirm(self, request, object, after_confirm, after_unsubscribe):
        conf = EmailConfirmation(
            content_object=object,
            after_confirm=after_confirm,
            after_unsubscribe=after_unsubscribe,
        )
        conf.send_email(request)

class EmailConfirmation(models.Model):
    confirmed = models.BooleanField(default=False)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')
    after_confirm = models.CharField(max_length=100)
    after_unsubscribe = models.CharField(max_length=100)

    objects = EmailConfirmationManager()

    def __unicode__(self):
        return 'Confirming of %s, %s' % (self.content_object.email, self.confirmed)

    def send_email(self, request):
        self.save()
        send_email(request, "Alert confirmation",
            'emailconfirmation/email.txt',
            {
                'object': self.content_object,
                'id': int_to_base32(self.id),
                'token': self.make_token(random.randint(0,32767)),
            }, self.content_object.email
        )

    @models.permalink
    def url_after_confirm(self):
        return (self.after_confirm, [ self.content_object.id ])

    @models.permalink
    def url_after_unsubscribe(self):
        return (self.after_unsubscribe, [ self.content_object.id ])

    def path_for_unsubscribe(self):
        return reverse('emailconfirmation-unsubscribe', kwargs={
            'id': int_to_base32(self.id),
            'token': self.make_token(random.randint(0,32767)),
        })

    def check_token(self, token):
        try:
            rand, hash = token.split("-")
        except:
            return False

        try:
            rand = base32_to_int(rand)
        except:
            return False

        if self.make_token(rand) != token:
            return False

        return True

    def make_token(self, rand):
        rand = int_to_base32(rand)
        hash = hmac.new(settings.SECRET_KEY, unicode(self.id) + rand, hashlib.sha1).hexdigest()[::2]
        return "%s-%s" % (rand, hash)
<|endoftext|>


<commit_before># Copyright (c) 2014-2015, Doug Kelly
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from django.db import models

from django.contrib.auth.models import User

# Create your models here.
class Registration(models.Model):
    user = models.ForeignKey(User)
    name = models.CharField(max_length=255)
    badge_name = models.CharField(max_length=32)
    email = models.EmailField()
    address = models.CharField(max_length=255)
    city = models.CharField(max_length=255)
    state = models.CharField(max_length=255)
    postal_code = models.CharField(max_length=255)
    country = models.CharField(max_length=255)
    birthday = models.DateField()
    registration_level = models.ForeignKey('RegistrationLevel')
    dealer_registration_level = models.ForeignKey('DealerRegistrationLevel', blank=True, null=True)
    shirt_size = models.ForeignKey('ShirtSize')
    volunteer = models.BooleanField(default=False, verbose_name='Contact me for volunteering opportunities')
    volunteer_phone = models.CharField(max_length=20, blank=True, verbose_name='Phone Number (only required if volunteering)')
    ip = models.GenericIPAddressField()

    def __unicode__(self):
        return self.name + ' [' + self.badge_name + ']'

class ShirtSize(models.Model):
    seq = models.IntegerField()
    size = models.CharField(max_length=20)

    def __unicode__(self):
        return self.size

class Payment(models.Model):
    registration = models.ForeignKey('Registration')
    payment_method = models.ForeignKey('PaymentMethod')
    payment_amount = models.DecimalField(max_digits=10, decimal_places=2)
    payment_received = models.DateTimeField(auto_now_add=True)

class RegistrationLevel(models.Model):
    seq = models.IntegerField()
    convention = models.ForeignKey('Convention')
    limit = models.IntegerField()
    title = models.CharField(max_length=255)
    price = models.DecimalField(max_digits=10, decimal_places=2)
    description = models.TextField()
    background = models.CharField(max_length=255)
    color = models.CharField(max_length=7)
    deadline = models.DateTimeField()
    active = models.BooleanField(default=True)

    def __unicode__(self):
        return self.title + ' [' + "%.02f" % (self.price) + ']'

class DealerRegistrationLevel(models.Model):
    convention = models.ForeignKey('Convention')
    number_tables = models.IntegerField()
    price = models.DecimalField(max_digits=10, decimal_places=2)

    def __unicode__(self):
        return str(self.number_tables) + ' [' + "%.02f" % (self.price) + ']'

class CouponCode(models.Model):
    convention = models.ForeignKey('Convention')
    code = models.CharField(max_length=255)
    discount = models.DecimalField(max_digits=10, decimal_places=2)
    percent = models.BooleanField(default=False)
    single_use = models.BooleanField(default=False)

    def __unicode__(self):
        if self.percent:
            return self.code + ' [' + str(self.discount) + '%]'
        else:
            return self.code + ' [' + "%.02f" % (self.discount) + ']'

class CouponUse(models.Model):
    registration = models.ForeignKey('Registration')
    coupon = models.ForeignKey('CouponCode')

class Convention(models.Model):
    name = models.CharField(max_length=255)
    description = models.TextField()
    agreement = models.TextField()
    dealer_limit = models.IntegerField()
    stripe_publishable_key = models.CharField(max_length=255)
    stripe_secret_key = models.CharField(max_length=255)

    def __unicode__(self):
        return self.name

class PaymentMethod(models.Model):
    seq = models.IntegerField()
    name = models.CharField(max_length=255)
    active = models.BooleanField(default=True)
    is_credit = models.BooleanField(default=False)

    def __unicode__(self):
        return self.name
<commit_msg>Add Python 3.x support for models<commit_after># Copyright (c) 2014-2015, Doug Kelly
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from django.db import models

from django.contrib.auth.models import User
from django.utils.encoding import python_2_unicode_compatible

# Create your models here.
@python_2_unicode_compatible
class Registration(models.Model):
    user = models.ForeignKey(User)
    name = models.CharField(max_length=255)
    badge_name = models.CharField(max_length=32)
    email = models.EmailField()
    address = models.CharField(max_length=255)
    city = models.CharField(max_length=255)
    state = models.CharField(max_length=255)
    postal_code = models.CharField(max_length=255)
    country = models.CharField(max_length=255)
    birthday = models.DateField()
    registration_level = models.ForeignKey('RegistrationLevel')
    dealer_registration_level = models.ForeignKey('DealerRegistrationLevel', verbose_name='Dealer Tables', blank=True, null=True)
    shirt_size = models.ForeignKey('ShirtSize')
    volunteer = models.BooleanField(default=False, verbose_name='Contact me for volunteering opportunities')
    volunteer_phone = models.CharField(max_length=20, blank=True, verbose_name='Phone Number (only required if volunteering)')
    ip = models.GenericIPAddressField()

    def __str__(self):
        return self.name + ' [' + self.badge_name + ']'

@python_2_unicode_compatible
class ShirtSize(models.Model):
    seq = models.IntegerField()
    size = models.CharField(max_length=20)

    def __str__(self):
        return self.size

class Payment(models.Model):
    registration = models.ForeignKey('Registration')
    payment_method = models.ForeignKey('PaymentMethod')
    payment_amount = models.DecimalField(max_digits=10, decimal_places=2)
    payment_received = models.DateTimeField(auto_now_add=True)

@python_2_unicode_compatible
class RegistrationLevel(models.Model):
    seq = models.IntegerField()
    convention = models.ForeignKey('Convention')
    limit = models.IntegerField()
    title = models.CharField(max_length=255)
    price = models.DecimalField(max_digits=10, decimal_places=2)
    description = models.TextField()
    background = models.CharField(max_length=255)
    color = models.CharField(max_length=7)
    deadline = models.DateTimeField()
    active = models.BooleanField(default=True)

    def __str__(self):
        return self.title + ' [' + "%.02f" % (self.price) + ']'

@python_2_unicode_compatible
class DealerRegistrationLevel(models.Model):
    convention = models.ForeignKey('Convention')
    number_tables = models.IntegerField()
    price = models.DecimalField(max_digits=10, decimal_places=2)

    def __str__(self):
        return str(self.number_tables) + ' [' + "%.02f" % (self.price) + ']'

@python_2_unicode_compatible
class CouponCode(models.Model):
    convention = models.ForeignKey('Convention')
    code = models.CharField(max_length=255)
    discount = models.DecimalField(max_digits=10, decimal_places=2)
    percent = models.BooleanField(default=False)
    single_use = models.BooleanField(default=False)

    def __str__(self):
        if self.percent:
            return self.code + ' [' + "%d%%" % (self.discount) + '%]'
        else:
            return self.code + ' [' + "%.02f" % (self.discount) + ']'

class CouponUse(models.Model):
    registration = models.ForeignKey('Registration')
    coupon = models.ForeignKey('CouponCode')

@python_2_unicode_compatible
class Convention(models.Model):
    name = models.CharField(max_length=255)
    description = models.TextField()
    agreement = models.TextField()
    dealer_limit = models.IntegerField()
    stripe_publishable_key = models.CharField(max_length=255)
    stripe_secret_key = models.CharField(max_length=255)

    def __str__(self):
        return self.name

@python_2_unicode_compatible
class PaymentMethod(models.Model):
    seq = models.IntegerField()
    name = models.CharField(max_length=255)
    active = models.BooleanField(default=True)
    is_credit = models.BooleanField(default=False)

    def __str__(self):
        return self.name
<|endoftext|>