<commit_before># function that bins and will be mapped in parallel
def binShapes(shapes, bBoxes, minBox, binWidth, bucketMin, ids = []):
    if not ids:
        ids = range(len(shapes))
    nShapes = len(ids)
    poly2Columns = dict([ (ids[i],set()) for i in range(nShapes) ])
    poly2Rows = dict([ (ids[i],set()) for i in range(nShapes) ])
    columns = dict([ (i,set()) for i in range(bucketMin + 2) ])
    rows = dict([ (i,set()) for i in range(bucketMin + 2)])

    for i in range(nShapes):
        idI = ids[i]
        shapeI = shapes[i]
        bBoxI = bBoxes[i][:]
        projBox = [int((bBoxI[j] - minBox[j]) / binWidth[j]) for j in range(4)]
        for j in range(projBox[0], projBox[2] + 1  ):
            columns[j].add(idI)
            poly2Columns[idI].add(j)
        for j in range(projBox[1], projBox[3] + 1):
            rows[j].add(idI)
            poly2Rows[idI].add(j)
    results = {}
    results['poly2Columns'] = poly2Columns
    results['poly2Rows'] = poly2Rows
    results['columns' ] = columns
    results['rows' ] = rows
    return results

def check_contiguity(shapes, bBoxes, poly2Rows, poly2Columns, rows, columns, ids = []):
    # note that ids is a list of particular shapes to check but all other arguments are not sliced

    def bbcommon(bb, bbother):
        """
        Checks for overlaps of bounding boxes. First, east-west, then north-south.
        Element 0 is west, element 2 is east, element 3 is north, element 1 is
        south.
        All four checks must be false for chflag to be true, meaning the two
        bounding boxes do overlap.
        """
        chflag = 0
        if not ((bbother[2] < bb[0]) or (bbother[0] > bb[2])):
            if not ((bbother[3] < bb[1]) or (bbother[1] > bb[3])):
                chflag = 1
        return chflag

    def queen(shapeA, shapeB):
        """
        Check if shapeA and shapeB are queen neighbors

        Arguments
        =========

        shapeA: pysal polygon object

        shapeB: pysal polygon object

        Returns
        =======

        1 if true, 0 if false

        Examples
        ========

        >>> sf = pysal.open(pysal.examples.get_path("columbus.shp"))
        >>> p0 = sf.get(0)
        >>> p1 = sf.get(1)
        >>> p3 = sf.get(3)
        >>> import binning
        >>> binning.queen(p0,p1)
        1
        >>> binning.queen(p0,p3)
        0
        >>> binning.queen(p1,p3)
        1

        """
        if bbcommon(shapeA.bounding_box, shapeB.bounding_box):
            a = set(shapeA.vertices)
            b = set(shapeB.vertices)
            common = a.intersection(b)
            if len(common) > 0:
                return 1
            else:
                return 0
        else:
            return 0

    if not ids:
        ids = poly2Rows.keys()

    neighbors = {}

    for polyId in ids:
        idRows = poly2Rows[polyId]
        idColumns = poly2Columns[polyId]
        rNeighbors = set()
        cNeighbors = set()
        for row in idRows:
            rNeighbors = rNeighbors.union(rows[row])
        for column in idColumns:
            cNeighbors = cNeighbors.union(columns[column])
        neighborsPolyId = rNeighbors.intersection(cNeighbors)
        if polyId not in neighbors:
            neighbors[polyId] = set()
        for j in neighborsPolyId:
            if polyId < j:
                if queen(shapes[polyId], shapes[j]):
                    if j not in neighbors:
                        neighbors[j] = set()
                    neighbors[j].add(polyId)
                    neighbors[polyId].add(j)
    return neighbors

if __name__ == '__main__':

    import pysal as ps
    import numpy as np
    import multiprocessing as mp
    import time

    sf = ps.examples.get_path("columbus.shp")
    shpFileObject = ps.open(sf)
    shapeBox = shpFileObject.bbox
    nShapes = len(shpFileObject)
    shapes = []
    bBoxes = []
    for shape in shpFileObject:
        shapes.append(shape)
        bBoxes.append(shape.bounding_box[:])

    # figure out grid that will be used by all processes
    DELTA = 0.000001
    bucketMin = nShapes / 80 + 2
    lengthY = ((shapeBox[2] + DELTA) - shapeBox[0]) / bucketMin
    lengthX = ((shapeBox[3] + DELTA) - shapeBox[1]) / bucketMin
    minBox = shapeBox[:2] * 2 # [minx, miny, minx, miny]
    binWidth = [ lengthX, lengthY] * 2 # [lenX, lenY, lenX, lenY]

    # sequential
    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)
    stage2Seq = check_contiguity(shapes, bBoxes, stage1Seq['poly2Rows'],
            stage1Seq['poly2Columns'], stage1Seq['rows'], stage1Seq['columns'])

    wseq = ps.W(stage2Seq)

    # parallel

    cores = mp.cpu_count()
    pool = mp.Pool(cores)
    print cores

    step = nShapes / (cores - 1)
    start = range(0, nShapes, step)[0:-1]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start, end)
    stage1 = {}
    for c in range(cores-1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage1[c] = pool.apply_async(binShapes, args=(shapes, bBoxes, minBox,
            binWidth, bucketMin, pids))
    pool.close()
    pool.join()

    results1 = {}
    for c in stage1:
        results1[c] = stage1[c].get()

    # combine results from stage 1

    columns = dict([ (i,set()) for i in range(bucketMin + 2) ])
    rows = dict([ (i,set()) for i in range(bucketMin + 2)])
    poly2Rows = {}
    poly2Columns = {}
    for c in  results1:
        result = results1[c]
        for row in rows:
            rows[row] = rows[row].union(result['rows'][row])
        for column in columns:
            columns[column] = columns[column].union(result['columns'][column])
        for polyId in result['poly2Columns']:
            poly2Columns[polyId] = result['poly2Columns'][polyId]
            poly2Rows[polyId] = result['poly2Rows'][polyId]

    # stage two
    pool = mp.Pool(cores)
    cuts = np.cumsum(np.arange(nShapes-1, 0, -1))
    cuts = cuts / (cuts[-1] / (cores - 1))
    start = [ np.nonzero(cuts==c)[0][0] for c in range((cores-1)) ]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start,end)
    print slices

    stage2 = {}
    for c in range(cores - 1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage2[c] = pool.apply_async(check_contiguity, args=(shapes, bBoxes,
            poly2Rows, poly2Columns, rows, columns, pids))
    pool.close()
    pool.join()
    results = {}
    for c in stage2:
        results[c] = stage2[c].get()

    # sequential step 1, parallel step 2

    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)

    pool = mp.Pool(cores)
    cuts = np.cumsum(np.arange(nShapes-1, 0, -1))
    cuts = cuts / (cuts[-1] / (cores - 1))
    start = [ np.nonzero(cuts==c)[0][0] for c in range((cores-1)) ]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start,end)
    print slices

    rows = stage1Seq['rows']
    columns = stage1Seq['columns']
    poly2Rows = stage1Seq['poly2Rows']
    poly2Columns = stage1Seq['poly2Columns']
    stage2 = {}
    for c in range(cores - 1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage2[c] = pool.apply_async(check_contiguity, args=(shapes, bBoxes,
            poly2Rows, poly2Columns, rows, columns, pids))
    pool.close()
    pool.join()
    results2 = {}
    for c in stage2:
        results2[c] = stage2[c].get()

    neighbors = dict([(i,set()) for i in range(nShapes) ])

    for c in results2:
        for key in results2[c]:
            neighbors[key] = neighbors[key].union(results2[c][key])

    wp2 = ps.W(neighbors)

    if wp2.neighbors == wseq.neighbors:
        print 'Parallel == Sequential'
    else:
        print 'Parallel != Sequential'

<commit_msg>- matching pysal with 1st stage sequential, 2nd stage parallel - next try parallelizing stage 1<commit_after># function that bins and will be mapped in parallel
def binShapes(shapes, bBoxes, minBox, binWidth, bucketMin, ids = []):
    if not ids:
        ids = range(len(shapes))
    nShapes = len(ids)
    poly2Columns = dict([ (ids[i],set()) for i in range(nShapes) ])
    poly2Rows = dict([ (ids[i],set()) for i in range(nShapes) ])
    columns = dict([ (i,set()) for i in range(bucketMin + 2) ])
    rows = dict([ (i,set()) for i in range(bucketMin + 2)])

    for i in range(nShapes):
        idI = ids[i]
        shapeI = shapes[i]
        bBoxI = bBoxes[i][:]
        projBox = [int((bBoxI[j] - minBox[j]) / binWidth[j]) for j in range(4)]
        #print i, projBox, bucketMin
        for j in range(projBox[0], projBox[2] + 1  ):
            columns[j].add(idI)
            poly2Columns[idI].add(j)
        for j in range(projBox[1], projBox[3] + 1):
            rows[j].add(idI)
            poly2Rows[idI].add(j)
    results = {}
    results['poly2Columns'] = poly2Columns
    results['poly2Rows'] = poly2Rows
    results['columns' ] = columns
    results['rows' ] = rows
    return results

def check_contiguity(shapes, bBoxes, poly2Rows, poly2Columns, rows, columns, ids = []):
    # note that ids is a list of particular shapes to check but all other arguments are not sliced

    def bbcommon(bb, bbother):
        """
        Checks for overlaps of bounding boxes. First, east-west, then north-south.
        Element 0 is west, element 2 is east, element 3 is north, element 1 is
        south.
        All four checks must be false for chflag to be true, meaning the two
        bounding boxes do overlap.
        """
        chflag = 0
        if not ((bbother[2] < bb[0]) or (bbother[0] > bb[2])):
            if not ((bbother[3] < bb[1]) or (bbother[1] > bb[3])):
                chflag = 1
        return chflag

    def queen(shapeA, shapeB):
        """
        Check if shapeA and shapeB are queen neighbors

        Arguments
        =========

        shapeA: pysal polygon object

        shapeB: pysal polygon object

        Returns
        =======

        1 if true, 0 if false

        Examples
        ========

        >>> sf = pysal.open(pysal.examples.get_path("columbus.shp"))
        >>> p0 = sf.get(0)
        >>> p1 = sf.get(1)
        >>> p3 = sf.get(3)
        >>> import binning
        >>> binning.queen(p0,p1)
        1
        >>> binning.queen(p0,p3)
        0
        >>> binning.queen(p1,p3)
        1

        """
        if bbcommon(shapeA.bounding_box, shapeB.bounding_box):
            a = set(shapeA.vertices)
            b = set(shapeB.vertices)
            common = a.intersection(b)
            if len(common) > 0:
                return 1
            else:
                return 0
        else:
            return 0

    if not ids:
        ids = poly2Rows.keys()

    neighbors = {}

    for polyId in ids:
        idRows = poly2Rows[polyId]
        idColumns = poly2Columns[polyId]
        rNeighbors = set()
        cNeighbors = set()
        for row in idRows:
            rNeighbors = rNeighbors.union(rows[row])
        for column in idColumns:
            cNeighbors = cNeighbors.union(columns[column])
        neighborsPolyId = rNeighbors.intersection(cNeighbors)
        if polyId not in neighbors:
            neighbors[polyId] = set()
        for j in neighborsPolyId:
            if polyId < j:
                if queen(shapes[polyId], shapes[j]):
                    if j not in neighbors:
                        neighbors[j] = set()
                    neighbors[j].add(polyId)
                    neighbors[polyId].add(j)
    return neighbors

if __name__ == '__main__':

    import pysal as ps
    import numpy as np
    import multiprocessing as mp
    import time

    """
    sf = ps.examples.get_path("columbus.shp")
    shpFileObject = ps.open(sf)
    shapeBox = shpFileObject.bbox
    nShapes = len(shpFileObject)
    shapes = []
    bBoxes = []
    for shape in shpFileObject:
        shapes.append(shape)
        bBoxes.append(shape.bounding_box[:])

    # figure out grid that will be used by all processes
    DELTA = 0.000001
    bucketMin = nShapes / 80 + 2
    lengthY = ((shapeBox[2] + DELTA) - shapeBox[0]) / bucketMin
    lengthX = ((shapeBox[3] + DELTA) - shapeBox[1]) / bucketMin
    minBox = shapeBox[:2] * 2 # [minx, miny, minx, miny]
    binWidth = [ lengthX, lengthY] * 2 # [lenX, lenY, lenX, lenY]

    # sequential
    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)
    stage2Seq = check_contiguity(shapes, bBoxes, stage1Seq['poly2Rows'],
            stage1Seq['poly2Columns'], stage1Seq['rows'], stage1Seq['columns'])

    wseq = ps.W(stage2Seq)

    # parallel

    cores = mp.cpu_count()
    pool = mp.Pool(cores)
    print cores

    step = nShapes / (cores - 1)
    start = range(0, nShapes, step)[0:-1]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start, end)
    stage1 = {}
    for c in range(cores-1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage1[c] = pool.apply_async(binShapes, args=(shapes, bBoxes, minBox,
            binWidth, bucketMin, pids))
    pool.close()
    pool.join()

    results1 = {}
    for c in stage1:
        results1[c] = stage1[c].get()

    # combine results from stage 1

    columns = dict([ (i,set()) for i in range(bucketMin + 2) ])
    rows = dict([ (i,set()) for i in range(bucketMin + 2)])
    poly2Rows = {}
    poly2Columns = {}
    for c in  results1:
        result = results1[c]
        for row in rows:
            rows[row] = rows[row].union(result['rows'][row])
        for column in columns:
            columns[column] = columns[column].union(result['columns'][column])
        for polyId in result['poly2Columns']:
            poly2Columns[polyId] = result['poly2Columns'][polyId]
            poly2Rows[polyId] = result['poly2Rows'][polyId]

    # stage two
    pool = mp.Pool(cores)
    cuts = np.cumsum(np.arange(nShapes-1, 0, -1))
    cuts = cuts / (cuts[-1] / (cores - 1))
    start = [ np.nonzero(cuts==c)[0][0] for c in range((cores-1)) ]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start,end)
    print slices

    stage2 = {}
    for c in range(cores - 1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage2[c] = pool.apply_async(check_contiguity, args=(shapes, bBoxes,
            poly2Rows, poly2Columns, rows, columns, pids))
    pool.close()
    pool.join()
    results = {}
    for c in stage2:
        results[c] = stage2[c].get()

    # sequential step 1, parallel step 2

    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)

    pool = mp.Pool(cores)
    cuts = np.cumsum(np.arange(nShapes-1, 0, -1))
    cuts = cuts / (cuts[-1] / (cores - 1))
    start = [ np.nonzero(cuts==c)[0][0] for c in range((cores-1)) ]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start,end)
    print slices

    rows = stage1Seq['rows']
    columns = stage1Seq['columns']
    poly2Rows = stage1Seq['poly2Rows']
    poly2Columns = stage1Seq['poly2Columns']
    stage2 = {}
    for c in range(cores - 1):
        pids = range(slices[c][0], slices[c][1])
        print pids
        stage2[c] = pool.apply_async(check_contiguity, args=(shapes, bBoxes,
            poly2Rows, poly2Columns, rows, columns, pids))
    pool.close()
    pool.join()
    results2 = {}
    for c in stage2:
        results2[c] = stage2[c].get()

    neighbors = dict([(i,set()) for i in range(nShapes) ])

    for c in results2:
        for key in results2[c]:
            neighbors[key] = neighbors[key].union(results2[c][key])

    wp2 = ps.W(neighbors)

    if wp2.neighbors == wseq.neighbors:
        print 'Parallel == Sequential'
    else:
        print 'Parallel != Sequential'
    """

    sfName = "nat.shp"

    t1 = time.time()
    sf = ps.examples.get_path(sfName)
    shpFileObject = ps.open(sf)
    shapeBox = shpFileObject.bbox
    nShapes = len(shpFileObject)
    shapes = []
    bBoxes = []
    for shape in shpFileObject:
        shapes.append(shape)
        bBoxes.append(shape.bounding_box[:])

    # figure out grid that will be used by all processes
    DELTA = 0.000001
    bucketMin = nShapes / 80 + 2
    lengthX = ((shapeBox[2] + DELTA) - shapeBox[0]) / bucketMin
    lengthY = ((shapeBox[3] + DELTA) - shapeBox[1]) / bucketMin
    minBox = shapeBox[:2] * 2 # [minx, miny, minx, miny]
    binWidth = [ lengthX, lengthY] * 2 # [lenX, lenY, lenX, lenY]

    # sequential
    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)
    stage2Seq = check_contiguity(shapes, bBoxes, stage1Seq['poly2Rows'],
            stage1Seq['poly2Columns'], stage1Seq['rows'], stage1Seq['columns'])

    wseq = ps.W(stage2Seq)
    t2 = time.time()
    print 'Sequential time: ', t2 - t1

    t1 = time.time()
    sf = ps.examples.get_path(sfName)
    shpFileObject = ps.open(sf)
    shapeBox = shpFileObject.bbox
    nShapes = len(shpFileObject)
    shapes = []
    bBoxes = []
    for shape in shpFileObject:
        shapes.append(shape)
        bBoxes.append(shape.bounding_box[:])

    # figure out grid that will be used by all processes
    DELTA = 0.000001
    bucketMin = nShapes / 80 + 2
    lengthX = ((shapeBox[2] + DELTA) - shapeBox[0]) / bucketMin
    lengthY = ((shapeBox[3] + DELTA) - shapeBox[1]) / bucketMin
    minBox = shapeBox[:2] * 2 # [minx, miny, minx, miny]
    binWidth = [ lengthX, lengthY] * 2 # [lenX, lenY, lenX, lenY]

    stage1Seq = binShapes(shapes, bBoxes, minBox, binWidth, bucketMin)

    cores = mp.cpu_count() - 1
    pool = mp.Pool(cores)
    cuts = np.cumsum(np.arange(nShapes-1, 0, -1))
    cuts = cuts / (cuts[-1] / (cores - 1))
    start = [ np.nonzero(cuts==c)[0][0] for c in range((cores-1)) ]
    end = start[1:]
    end.append(nShapes)
    slices = zip(start,end)

    rows = stage1Seq['rows']
    columns = stage1Seq['columns']
    poly2Rows = stage1Seq['poly2Rows']
    poly2Columns = stage1Seq['poly2Columns']
    stage2 = {}
    for c in range(cores - 1):
        pids = range(slices[c][0], slices[c][1])
        stage2[c] = pool.apply_async(check_contiguity, args=(shapes, bBoxes,
            poly2Rows, poly2Columns, rows, columns, pids))
    pool.close()
    pool.join()
    results2 = {}
    for c in stage2:
        results2[c] = stage2[c].get()

    neighbors = dict([(i,set()) for i in range(nShapes) ])

    for c in results2:
        for key in results2[c]:
            neighbors[key] = neighbors[key].union(results2[c][key])

    wp2 = ps.W(neighbors)
    t2 = time.time()
    print 'Parallel (seq stage 1, par stage 2)', t2 - t1

    if wp2.neighbors == wseq.neighbors:
        print 'Same weights'
    else:
        print 'different weights'

    t1 = time.time()
    wpysal = ps.queen_from_shapefile(sf)
    t2 = time.time()
    print 'Sequential pysal: ', t2 - t1
    nBad = 0
    for key in wpysal.neighbors:
        if set(wpysal.neighbors[key]) != wp2.neighbors[key]:
                nBad += 1
    if nBad == 0:
        print 'Same weights as pysal: ', True
    else:
        print 'Same weights as pysal: ', False

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import sys
import json
import yaml
import argparse
from time import time

from config import validate_config
from utils import usage, java_string_hashcode
from row_generator import RowGenerator
from kafka import KafkaProducer

UPDATE_PERIOD_MILLIS = 250
DISPLAY_PROGRESS_WARMUP_MILLIS = 5000

def display_line(line):
    ansiClearLine = "\033[K"
    ansiMoveCursorToColumnZero = "\033[0G"
    sys.stdout.write(ansiClearLine + ansiMoveCursorToColumnZero + line)
    sys.stdout.flush()

def display_progress(total, count, started_time_millis):
    current_time_millis = time() * 1000.0
    elapsed_millis = current_time_millis - started_time_millis
    predicted_total_millis = (elapsed_millis / count) * total
    remaining_millis = predicted_total_millis - elapsed_millis
    duration = pretty_duration(remaining_millis, elapsed_millis)
    display_line("%d / %d (%.2f%%) %s" % (count, total, ( count * 100.0 ) / total, duration))

def pretty_duration(millis, elapsedMillis):
    if elapsedMillis < DISPLAY_PROGRESS_WARMUP_MILLIS:
        return ""
    d = (millis / (1000 * 60 * 60 * 24))
    h = (millis / (1000 * 60 * 60)) % 24
    m = (millis / (1000 * 60)) % 60
    s = (millis / (1000)) % 60
    if d > 0:
        return "- %d days %02dh %02dm %02ds remaining " % (d, h, m, s)
    elif h > 0:
        return "- %02dh %02dm %02ds remaining " % (h, m, s)
    elif m > 0:
        return "- %02dm %02ds remaining " % (m, s)
    elif s > 0:
        return "- %02ds remaining " % s
    else:
        return ""

def produce(producer, topic, partition, message):
    producer.send(topic, key = "hello", value = message, partition = partition)

def maxwell_message(database, table, type, data):
    return json.dumps({
        "database": database,
        "table": table,
        "type": type,
        "ts": int(time()),
        "data": data
    })

def bootstrap_start_message(database, table):
    return maxwell_message(database, table, "bootstrap-start", {})

def bootstrap_complete_message(database, table):
    return maxwell_message(database, table, "bootstrap-complete", {})

def bootstrap(producer, schema, database, table, config):
    topic = config['kafka']['topic']
    start_time_millis = time() * 1000.0
    inserted_rows = 0
    total_rows = int(float(config['mysql']['schemas'][schema]['tables'][table][database]['bootstrap-count']))
    partition_count = 1 + max(producer.partitions_for(topic))
    partition = abs(java_string_hashcode(database) % partition_count)
    produce(producer, topic, partition, bootstrap_start_message(database, table))
    for message in bootstrap_insert_messages(schema, database, table, config, total_rows):
        produce(producer, topic, partition, message)
        inserted_rows += 1
        display_progress(total_rows, inserted_rows, start_time_millis)
    produce(producer, topic, partition, bootstrap_complete_message(database, table))
    display_line("")

def bootstrap_insert_messages(schema, database, table, config, rows_total):
    row_generator = RowGenerator.get_instance(schema, database, table, config)
    for row_index in xrange(rows_total):
        yield maxwell_message(database, table, "bootstrap-insert", row_generator.generate_row(row_index))

def find_schema(config, database, table):
    found_schema = None
    for schema in config['mysql']['schemas']:
        databases = config['mysql']['schemas'][schema]['databases']
        tables = config['mysql']['schemas'][schema]['tables']
        if database in databases and table in tables:
            if found_schema: usage('two different schemas cannot have identical database names')
            found_schema = schema
    if not found_schema: usage('could not find schema with specified database and table')
    return found_schema

def main():
    parser = argparse.ArgumentParser(description='Maxwell faker for systems and load testing.')
    parser.add_argument('--config', metavar='CONFIG', type=str, required=True, help='path to yaml config file')
    parser.add_argument('--database', metavar='DATABASE', type=str, required=True, help='database to bootstrap')
    parser.add_argument('--table', metavar='TABLE', type=str, required=True, help='table to bootstrap')
    args = parser.parse_args()
    config = yaml.load(open(args.config).read())
    validate_config(config)
    schema = find_schema(config, args.database, args.table)
    producer = KafkaProducer(bootstrap_servers = config['kafka']['brokers'])
    try:
        bootstrap(producer, schema, args.database, args.table, config)
    except IOError, e:
        usage(e)
    except KeyboardInterrupt:
        print
    finally:
        print "flushing producer..." + 20 * " "
        producer.flush()
        producer.close()
        print "done"
<commit_msg>Kafka producer now produces correct keys.<commit_after># -*- coding: utf-8 -*-

import sys
import json
import yaml
import argparse
from time import time

from config import validate_config
from utils import usage, java_string_hashcode
from row_generator import RowGenerator
from kafka import KafkaProducer

UPDATE_PERIOD_MILLIS = 250
DISPLAY_PROGRESS_WARMUP_MILLIS = 5000

def display_line(line):
    ansiClearLine = "\033[K"
    ansiMoveCursorToColumnZero = "\033[0G"
    sys.stdout.write(ansiClearLine + ansiMoveCursorToColumnZero + line)
    sys.stdout.flush()

def display_progress(total, count, started_time_millis):
    current_time_millis = time() * 1000.0
    elapsed_millis = current_time_millis - started_time_millis
    predicted_total_millis = (elapsed_millis / count) * total
    remaining_millis = predicted_total_millis - elapsed_millis
    duration = pretty_duration(remaining_millis, elapsed_millis)
    display_line("%d / %d (%.2f%%) %s" % (count, total, ( count * 100.0 ) / total, duration))

def pretty_duration(millis, elapsedMillis):
    if elapsedMillis < DISPLAY_PROGRESS_WARMUP_MILLIS:
        return ""
    d = (millis / (1000 * 60 * 60 * 24))
    h = (millis / (1000 * 60 * 60)) % 24
    m = (millis / (1000 * 60)) % 60
    s = (millis / (1000)) % 60
    if d > 0:
        return "- %d days %02dh %02dm %02ds remaining " % (d, h, m, s)
    elif h > 0:
        return "- %02dh %02dm %02ds remaining " % (h, m, s)
    elif m > 0:
        return "- %02dm %02ds remaining " % (m, s)
    elif s > 0:
        return "- %02ds remaining " % s
    else:
        return ""

def produce(producer, topic, partition, key, value):
    producer.send(topic, key = key, value = value, partition = partition)

def maxwell_message(database, table, type, data, pk_name, pk_value):
    key = json.dumps({
        "database": database,
        "table": table,
        "pk." + pk_name: pk_value
    })
    value = json.dumps({
        "database": database,
        "table": table,
        "type": type,
        "ts": int(time()),
        "data": data
    })
    return key, value

def bootstrap_start_message(schema, database, table, config):
    row_generator = RowGenerator.get_instance(schema, database, table, config)
    data = {}
    pk_name = row_generator.primary_key_field
    pk_value = None
    return maxwell_message(database, table, "bootstrap-start", data, pk_name, pk_value)

def bootstrap_complete_message(schema, database, table, config):
    row_generator = RowGenerator.get_instance(schema, database, table, config)
    data = {}
    pk_name = row_generator.primary_key_field
    pk_value = None
    return maxwell_message(database, table, "bootstrap-complete", data, pk_name, pk_value)

def bootstrap(producer, schema, database, table, config):
    topic = config['kafka']['topic']
    start_time_millis = time() * 1000.0
    inserted_rows = 0
    total_rows = int(float(config['mysql']['schemas'][schema]['tables'][table][database]['bootstrap-count']))
    partition_count = 1 + max(producer.partitions_for(topic))
    partition = abs(java_string_hashcode(database) % partition_count)
    produce(producer, topic, partition, *bootstrap_start_message(schema, database, table, config))
    for key, value in bootstrap_insert_messages(schema, database, table, config, total_rows):
        produce(producer, topic, partition, key, value)
        inserted_rows += 1
        display_progress(total_rows, inserted_rows, start_time_millis)
    produce(producer, topic, partition, *bootstrap_complete_message(schema, database, table, config))
    display_line("")

def bootstrap_insert_messages(schema, database, table, config, rows_total):
    row_generator = RowGenerator.get_instance(schema, database, table, config)
    for row_index in xrange(rows_total):
        data = row_generator.generate_row(row_index)
        pk_name = row_generator.primary_key_field
        pk_value = row_generator.generate_primary_key(row_index)
        yield maxwell_message(database, table, "bootstrap-insert", data, pk_name, pk_value)

def find_schema(config, database, table):
    found_schema = None
    for schema in config['mysql']['schemas']:
        databases = config['mysql']['schemas'][schema]['databases']
        tables = config['mysql']['schemas'][schema]['tables']
        if database in databases and table in tables:
            if found_schema: usage('two different schemas cannot have identical database names')
            found_schema = schema
    if not found_schema: usage('could not find schema with specified database and table')
    return found_schema

def main():
    parser = argparse.ArgumentParser(description='Maxwell faker for systems and load testing.')
    parser.add_argument('--config', metavar='CONFIG', type=str, required=True, help='path to yaml config file')
    parser.add_argument('--database', metavar='DATABASE', type=str, required=True, help='database to bootstrap')
    parser.add_argument('--table', metavar='TABLE', type=str, required=True, help='table to bootstrap')
    args = parser.parse_args()
    config = yaml.load(open(args.config).read())
    validate_config(config)
    schema = find_schema(config, args.database, args.table)
    producer = KafkaProducer(bootstrap_servers = config['kafka']['brokers'])
    try:
        bootstrap(producer, schema, args.database, args.table, config)
    except KeyboardInterrupt:
        sys.exit(1)
    producer.flush()
    producer.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf8 -*-

from .topping import Topping

from jawa.constants import *
from jawa.cf import ClassFile
from jawa.util.descriptor import method_descriptor, field_descriptor

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

class BlockStateTopping(Topping):
    """Gets tile entity (block entity) types."""

    PROVIDES = [
        "blocks.states"
    ]

    DEPENDS = [
        "blocks",
        "version.data",
        "identify.blockstatecontainer"
    ]

    @staticmethod
    def act(aggregate, jar, verbose=False):
        if "blockstatecontainer" not in aggregate["classes"]:
            if verbose:
                print "blockstatecontainer not found; skipping blockstates"
            return

        # 1449 is 17w46a
        is_flattened = ("data" in aggregate["version"] and aggregate["version"]["data"] > 1449)

        blockstatecontainer = aggregate["classes"]["blockstatecontainer"]
        block_cf = ClassFile(StringIO(jar.read(aggregate["classes"]["block.superclass"] + ".class")))
        plane = aggregate["classes"]["enumfacing.plane"]

        base_method = block_cf.methods.find_one(returns="L" + blockstatecontainer + ";", f=lambda m: m.access_flags.acc_protected)
        def matches(other_method):
            return (other_method.name.value == base_method.name.value and
                    other_method.descriptor.value == base_method.descriptor.value)

        _property_types = set()
        # Properties that are used by each block class
        properties_by_class = {}
        def process_class(name):
            """
            Gets the properties for the given block class, checking the parent
            class if none are defined.  Returns the properties, and also adds
            them to properties_by_class
            """
            if name in properties_by_class:
                # Caching - avoid reading the same class multiple times
                return properties_by_class[name]

            cf = ClassFile(StringIO(jar.read(name + ".class")))
            method = cf.methods.find_one(f=matches)

            if not method:
                properties = process_class(cf.super_.name.value)
                properties_by_class[name] = properties
                return properties

            properties = None
            stack = []
            for ins in method.code.disassemble():
                # This could _almost_ just be checking for getstatic, but
                # brewing stands use an array of properties as the field,
                # so we need some stupid extra logic.
                if ins.mnemonic == "new":
                    const = cf.constants.get(ins.operands[0].value)
                    type_name = const.name.value
                    assert type_name == blockstatecontainer
                    stack.append(object())
                elif ins.mnemonic.startswith("iconst"):
                    stack.append(int(ins.mnemonic[-1]))
                elif ins.mnemonic.endswith("ipush"):
                    stack.append(ins.operands[0].value)
                elif ins.mnemonic == "anewarray":
                    len = stack.pop()
                    val = [None] * len
                    if not properties:
                        properties = val
                    stack.append(val)
                elif ins.mnemonic == "getstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    prop = {
                        "field_name": const.name_and_type.name.value
                    }
                    desc = field_descriptor(const.name_and_type.descriptor.value)
                    _property_types.add(desc.name)
                    stack.append(prop)
                elif ins.mnemonic == "aaload":
                    index = stack.pop()
                    array = stack.pop()
                    prop = array.copy()
                    prop["array_index"] = index
                    stack.append(prop)
                elif ins.mnemonic == "aastore":
                    value = stack.pop()
                    index = stack.pop()
                    array = stack.pop()
                    array[index] = value
                elif ins.mnemonic == "dup":
                    stack.append(stack[-1])
                elif ins.mnemonic == "invokespecial":
                    break

            properties_by_class[name] = properties
            return properties

        for block in aggregate["blocks"]["block"].itervalues():
            process_class(block["class"])

        assert len(_property_types) == 4
        property_types = {}
        for type in _property_types:
            cf = ClassFile(StringIO(jar.read(type + ".class")))
            if cf.super_.name.value in _property_types:
                property_types[type] = "direction"
            else:
                attribute = cf.attributes.find_one(name='Signature')
                signature = attribute.signature.value
                # Somewhat ugly behavior until an actual parser is added for these
                if "Enum" in signature:
                    property_types[type] = "enum"
                elif "Integer" in signature:
                    property_types[type] = "int"
                elif "Boolean" in signature:
                    property_types[type] = "bool"
                elif verbose:
                    print "Unknown property type %s with signature %s" % (type, signature)

        print property_types
        def find_field(cls, field_name):
            cf = ClassFile(StringIO(jar.read(cls + ".class")))
            if not cf.fields.find_one(field_name):
                return find_field(cf.super_.name.value, field_name)

            init = cf.methods.find_one("<clinit>")
            stack = []
            for ins in init.code.disassemble():
                if ins.mnemonic == "putstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    value = stack.pop()
                    if const.name_and_type.name.value == field_name:
                        return value
                elif ins.mnemonic == "getstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    target = const.class_.name.value
                    type = field_descriptor(const.name_and_type.descriptor.value).name
                    name = const.name_and_type.name.value
                    cls2 = const.class_.name.value
                    if type in property_types:
                        stack.append(find_field(cls2, name))
                    else:
                        stack.append(find_field(target, name))
                elif ins.mnemonic in ("ldc", "ldc_w", "ldc2_w"):
                    const = cf.constants.get(ins.operands[0].value)

                    if isinstance(const, ConstantClass):
                        stack.append("%s.class" % const.name.value)
                    elif isinstance(const, ConstantString):
                        stack.append(const.string.value)
                    else:
                        stack.append(const.value)
                elif ins.mnemonic.startswith("iconst"):
                    stack.append(int(ins.mnemonic[-1]))
                elif ins.mnemonic.startswith("dconst"):
                    stack.append(float(ins.mnemonic[-1]))
                elif ins.mnemonic.endswith("ipush"):
                    stack.append(ins.operands[0].value)
                elif ins.mnemonic == "anewarray":
                    length = stack.pop()
                    stack.append([None] * length)
                elif ins.mnemonic == "aastore":
                    value = stack.pop()
                    index = stack.pop()
                    array = stack.pop()
                    array[index] = value
                elif ins.mnemonic == "dup":
                    stack.append(stack[-1])
                elif ins.mnemonic.startswith("invoke"):
                    const = cf.constants.get(ins.operands[0].value)
                    desc = method_descriptor(const.name_and_type.descriptor.value)
                    num_args = len(desc.args)
                    args = [stack.pop() for _ in xrange(num_args)]
                    args.reverse()

                    if ins.mnemonic == "invokestatic":
                        obj = None
                    else:
                        obj = stack.pop()

                    if desc.returns.name in property_types:
                        name = args[0]
                        prop = {
                            "name": name,
                            "class": desc.returns.name,
                            "type": property_types[desc.returns.name],
                            "args": args
                        }
                        stack.append(prop)
                    elif const.name_and_type.name.value == "<init>":
                        if obj["is_enum"]:
                            obj["enum_name"] = args[0]
                        else:
                            obj["args"] = args
                    elif desc.returns.name != "void":
                        if obj and obj["class"] == plane:
                            # One special case, where EnumFacing.Plane is used
                            # to get a list of directions
                            assert obj["enum_name"] in ("HORIZONTAL", "VERTICAL")
                            if obj["enum_name"] == "HORIZONTAL":
                                stack.append(["NORTH", "EAST", "SOUTH", "WEST"])
                            else:
                                stack.append(["UP", "DOWN"])
                        else:
                            stack.append(object())
                elif ins.mnemonic == "new":
                    const = cf.constants.get(ins.operands[0].value)
                    type_name = const.name.value
                    tcf = ClassFile(StringIO(jar.read(type_name + ".class")))
                    obj = {
                        "class": type_name,
                        "is_enum": tcf.super_.name.value == "java/lang/Enum"
                    }
                    stack.append(obj)
                else:
                    print "Unhandled:", cls, ins

        for cls, properties in properties_by_class.iteritems():
            #print cls, properties
            # TODO: Optimize this - less class loading!
            for property in properties:
                field_name = property["field_name"]
                try:
                    field = find_field(cls, field_name)
                    if "array_index" in property:
                        field = field[property["array_index"]]
                    print field
                    property["value"] = field
                except Exception as e:
                    print cls, property
                    print e

        for block in aggregate["blocks"]["block"].itervalues():
            block["states"] = [it["value"] for it in properties_by_class[block["class"]]]

<commit_msg>Cache class fields<commit_after>#!/usr/bin/env python
# -*- coding: utf8 -*-

from .topping import Topping

from jawa.constants import *
from jawa.cf import ClassFile
from jawa.util.descriptor import method_descriptor, field_descriptor

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

class BlockStateTopping(Topping):
    """Gets tile entity (block entity) types."""

    PROVIDES = [
        "blocks.states"
    ]

    DEPENDS = [
        "blocks",
        "version.data",
        "identify.blockstatecontainer"
    ]

    @staticmethod
    def act(aggregate, jar, verbose=False):
        if "blockstatecontainer" not in aggregate["classes"]:
            if verbose:
                print "blockstatecontainer not found; skipping blockstates"
            return

        # 1449 is 17w46a
        is_flattened = ("data" in aggregate["version"] and aggregate["version"]["data"] > 1449)

        blockstatecontainer = aggregate["classes"]["blockstatecontainer"]
        block_cf = ClassFile(StringIO(jar.read(aggregate["classes"]["block.superclass"] + ".class")))
        plane = aggregate["classes"]["enumfacing.plane"]

        base_method = block_cf.methods.find_one(returns="L" + blockstatecontainer + ";", f=lambda m: m.access_flags.acc_protected)
        def matches(other_method):
            return (other_method.name.value == base_method.name.value and
                    other_method.descriptor.value == base_method.descriptor.value)

        _property_types = set()
        # Properties that are used by each block class
        properties_by_class = {}
        def process_class(name):
            """
            Gets the properties for the given block class, checking the parent
            class if none are defined.  Returns the properties, and also adds
            them to properties_by_class
            """
            if name in properties_by_class:
                # Caching - avoid reading the same class multiple times
                return properties_by_class[name]

            cf = ClassFile(StringIO(jar.read(name + ".class")))
            method = cf.methods.find_one(f=matches)

            if not method:
                properties = process_class(cf.super_.name.value)
                properties_by_class[name] = properties
                return properties

            properties = None
            stack = []
            for ins in method.code.disassemble():
                # This could _almost_ just be checking for getstatic, but
                # brewing stands use an array of properties as the field,
                # so we need some stupid extra logic.
                if ins.mnemonic == "new":
                    const = cf.constants.get(ins.operands[0].value)
                    type_name = const.name.value
                    assert type_name == blockstatecontainer
                    stack.append(object())
                elif ins.mnemonic.startswith("iconst"):
                    stack.append(int(ins.mnemonic[-1]))
                elif ins.mnemonic.endswith("ipush"):
                    stack.append(ins.operands[0].value)
                elif ins.mnemonic == "anewarray":
                    len = stack.pop()
                    val = [None] * len
                    if not properties:
                        properties = val
                    stack.append(val)
                elif ins.mnemonic == "getstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    prop = {
                        "field_name": const.name_and_type.name.value
                    }
                    desc = field_descriptor(const.name_and_type.descriptor.value)
                    _property_types.add(desc.name)
                    stack.append(prop)
                elif ins.mnemonic == "aaload":
                    index = stack.pop()
                    array = stack.pop()
                    prop = array.copy()
                    prop["array_index"] = index
                    stack.append(prop)
                elif ins.mnemonic == "aastore":
                    value = stack.pop()
                    index = stack.pop()
                    array = stack.pop()
                    array[index] = value
                elif ins.mnemonic == "dup":
                    stack.append(stack[-1])
                elif ins.mnemonic == "invokespecial":
                    break

            properties_by_class[name] = properties
            return properties

        for block in aggregate["blocks"]["block"].itervalues():
            process_class(block["class"])

        assert len(_property_types) == 4
        property_types = {}
        for type in _property_types:
            cf = ClassFile(StringIO(jar.read(type + ".class")))
            if cf.super_.name.value in _property_types:
                property_types[type] = "direction"
            else:
                attribute = cf.attributes.find_one(name='Signature')
                signature = attribute.signature.value
                # Somewhat ugly behavior until an actual parser is added for these
                if "Enum" in signature:
                    property_types[type] = "enum"
                elif "Integer" in signature:
                    property_types[type] = "int"
                elif "Boolean" in signature:
                    property_types[type] = "bool"
                elif verbose:
                    print "Unknown property type %s with signature %s" % (type, signature)

        print property_types
        fields_by_class = {}
        ignore_scary_lambda_marker = object()
        def find_field(cls, field_name):
            """
            cls: name of the class
            field_name: name of the field to find.  If None, returns all fields
            """
            if cls in fields_by_class:
                if fields_by_class[cls] == ignore_scary_lambda_marker:
                    return object()
                elif field_name is not None:
                    return fields_by_class[cls][field_name]
                else:
                    return fields_by_class[cls]

            cf = ClassFile(StringIO(jar.read(cls + ".class")))

            if cf.attributes.find_one("BootstrapMethods"):
                # AAAH LAMBDAS
                fields_by_class[cls] = ignore_scary_lambda_marker
                return object()

            fields_by_class[cls] = {}
            super_name = cf.super_.name.value
            if not super_name.startswith("java/lang"):
                # Add fields from superclass
                fields_by_class[cls].update(find_field(super_name, None))

            init = cf.methods.find_one("<clinit>")
            if not init:
                if field_name is not None:
                    return fields_by_class[cls][field_name]
                else:
                    return fields_by_class[cls]

            stack = []
            for ins in init.code.disassemble():
                if ins.mnemonic == "putstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    name = const.name_and_type.name.value
                    value = stack.pop()

                    if isinstance(value, dict):
                        value["declared_in"] = cls
                    fields_by_class[cls][name] = value
                elif ins.mnemonic == "getstatic":
                    const = cf.constants.get(ins.operands[0].value)
                    target = const.class_.name.value
                    type = field_descriptor(const.name_and_type.descriptor.value).name
                    name = const.name_and_type.name.value
                    if not target.startswith("java/"):

                        try:
                            stack.append(find_field(target, name))
                        except:
                            print 'Failed to handle', target, name, 'for', cls
                            raise
                    else:
                        stack.append(object())
                elif ins.mnemonic in ("ldc", "ldc_w", "ldc2_w"):
                    const = cf.constants.get(ins.operands[0].value)

                    if isinstance(const, ConstantClass):
                        stack.append("%s.class" % const.name.value)
                    elif isinstance(const, ConstantString):
                        stack.append(const.string.value)
                    else:
                        stack.append(const.value)
                elif ins.mnemonic.startswith("iconst"):
                    stack.append(int(ins.mnemonic[-1]))
                elif ins.mnemonic.startswith("dconst"):
                    stack.append(float(ins.mnemonic[-1]))
                elif ins.mnemonic.endswith("ipush"):
                    stack.append(ins.operands[0].value)
                elif ins.mnemonic == "anewarray":
                    length = stack.pop()
                    stack.append([None] * length)
                elif ins.mnemonic == "aastore":
                    value = stack.pop()
                    index = stack.pop()
                    array = stack.pop()
                    array[index] = value
                elif ins.mnemonic == "dup":
                    stack.append(stack[-1])
                elif ins.mnemonic.startswith("invoke"):
                    const = cf.constants.get(ins.operands[0].value)
                    desc = method_descriptor(const.name_and_type.descriptor.value)
                    num_args = len(desc.args)
                    args = [stack.pop() for _ in xrange(num_args)]
                    args.reverse()

                    if ins.mnemonic == "invokestatic":
                        if const.class_.name.value.startswith("com/google/"):
                            # Call to e.g. Maps.newHashMap, beyond what we
                            # care about
                            break
                        obj = None
                    else:
                        obj = stack.pop()

                    if desc.returns.name in property_types:
                        name = args[0]
                        prop = {
                            "name": name,
                            "class": desc.returns.name,
                            "type": property_types[desc.returns.name],
                            "args": args
                        }
                        stack.append(prop)
                    elif const.name_and_type.name.value == "<init>":
                        if obj["is_enum"]:
                            obj["enum_name"] = args[0]
                        else:
                            obj["args"] = args
                    elif desc.returns.name != "void":
                        if isinstance(obj, dict) and obj["class"] == plane:
                            # One special case, where EnumFacing.Plane is used
                            # to get a list of directions
                            assert obj["enum_name"] in ("HORIZONTAL", "VERTICAL")
                            if obj["enum_name"] == "HORIZONTAL":
                                stack.append(["NORTH", "EAST", "SOUTH", "WEST"])
                            else:
                                stack.append(["UP", "DOWN"])
                        else:
                            stack.append(object())
                elif ins.mnemonic == "new":
                    const = cf.constants.get(ins.operands[0].value)
                    type_name = const.name.value
                    tcf = ClassFile(StringIO(jar.read(type_name + ".class")))
                    obj = {
                        "class": type_name,
                        "is_enum": tcf.super_.name.value == "java/lang/Enum"
                    }
                    stack.append(obj)
                elif ins.mnemonic == "return":
                    break
                elif ins.mnemonic == "if_icmpge":
                    # Code in stairs that loops over state combinations for hitboxes
                    break

            if field_name is not None:
                return fields_by_class[cls][field_name]
            else:
                return fields_by_class[cls]

        for cls, properties in properties_by_class.iteritems():
            #print cls, properties
            # TODO: Optimize this - less class loading!
            for property in properties:
                field_name = property["field_name"]
                try:
                    field = find_field(cls, field_name)
                    if "array_index" in property:
                        field = field[property["array_index"]]
                    #print field
                    property["field"] = field
                except Exception as e:
                    print cls + "." + field_name, property
                    import traceback
                    traceback.print_exc()
                    #print e

        #print fields_by_class

        for block in aggregate["blocks"]["block"].itervalues():
            block["states"] = [it["field"] for it in properties_by_class[block["class"]]]

<|endoftext|>


<commit_before>"""
    Import / Export of content

    Requirements:

    - exporting of new types should work, but can be customized
    - xml definitions are versioned, backwards compatibility should be
      present

    - should position be exported? Or should xml order be used, and position
      be renumbered when importing?

    - due to the current nature of nodes, content with a path of /foo/bar may
      exist without a node named /foo. This is an inconsistency in the node tree.
      Such an implementation detail is not relevant to the import/export,
      but how can we handle content without a strict parent? Make it unattached?

    - how to export nodes without content? (unattached nodes)

    - how to handle related content? This can be a User object (its id can
      change) or entirely new models (custom export?)

    Structure (roughly)
    <site version="1" ..>
      <content type=".." slug="..." position="1">
        <created></created>
        <publication></publication>
        <expires></expires>
        <title></title>

        <children>
          <content type=".." slug="..">
          </content>
        </children>
      </content>
      <unattached>

      </unattached>
    </site>

    Should a field be serialized to <title> or <field name="title">? The
    former will not allow for a schema definition
"""

from xml.etree.ElementTree import Element, SubElement, tostring
from django.utils.encoding import smart_unicode
from django.contrib.auth.models import User

from .content import type_registry

class SerializationException(Exception):
    pass

class WheelSerializer(object):
    """ (de) serialize wheel content """
    skip = ('node', )

    def __init__(self):
        pass

    def serialize_owner(self, field, o):
        # import pytest; pytest.set_trace()
        value = getattr(o, field.name, None)
        # value = smart_unicode(value)
        if value is None:
            return ""
        return smart_unicode(value.username)

    def deserialize_owner(self, field, tree):
        username = tree.text
        if not username:
            return None
        return User.objects.get(username=username)

    def serialize(self, spoke):
        o = spoke.instance
        fields = {}
        for field in o._meta.concrete_model._meta.fields:
            handler = getattr(self, "serialize_%s" % field.name, None)
            if field.name in self.skip:
                continue

            if handler:
                fields[field.name] = handler(field, o)
            elif field.serialize:
                if field.rel is None:
                    value = field.value_to_string(o)
                else:
                    value = getattr(o, field.get_attname())
                    value = smart_unicode(value)

                fields[field.name] = value
        xmlfields = Element("fields")
        for k, v in fields.iteritems():
            e = SubElement(xmlfields, "field")
            e.attrib['name'] = k
            e.text = v
        return xmlfields

    def deserialize(self, spoke, tree):
        model = spoke.model
        fields = {}
        # import pytest; pytest.set_trace()
        for field_node in tree.findall("field"):
            field_name = field_node.attrib.get("name")
            if not field_name:
                raise SerializationException("Missing name attribute")
            field = model._meta.get_field(field_name)

            handler = getattr(self, "deserialize_%s" % field.name, None)
            if handler:
                fields[field_name] = handler(field, field_node)

            ##elif field.rel and isinstance(field.rel, models.ManyToManyRel):
            ##    m2m_data[field.name] = self._handle_m2m_field_node(field_node, field)
            ##elif field.rel and isinstance(field.rel, models.ManyToOneRel):
            ##    data[field.attname] = self._handle_fk_field_node(field_node, field)
            else:
                value = field.to_python(field_node.text)
                fields[field.name] = value
        m = model(**fields).save()
        return spoke(m)

class Exporter(object):
    VERSION = 1

    def __init__(self, verbose=0):
        self.verbose = verbose

    def export_node(self, parent, node):
        # import pdb; pdb.set_trace()

        try:
            spoke = node.content().spoke()
        except AttributeError:
            spoke = None

        xmlcontent = SubElement(parent, "content",
                                dict(slug=node.slug(),
                                type=spoke.model.get_name()))
        if spoke:
            contentxml = spoke.serializer().serialize(spoke)
            xmlcontent.append(contentxml)

        children = SubElement(xmlcontent, "children")
        for child in node.children():
            self.export_node(children, child)

    def run(self, node, base="", unattached=True, attachinline=True):
        """ export node and all content beneath it. If unattached is True,
            also export unattached content.

            attachments can be inline or exported to a folder
        """
        root = Element("site")
        root.set('version', str(self.VERSION))
        root.set('base', base)
        self.export_node(root, node)
        return root

class Importer(object):
    def __init__(self, verbose=0):
        self.verbose = verbose

    def import_node(self, node, tree):
        type = tree.attrib['type']
        slug = tree.attrib['slug']
        spoke = type_registry.get(type)
        s = spoke.serializer().deserialize(spoke, tree)
        if slug == "":
            n = node
        else:
            n = node.add(slug)
        n.set(s.instance)

        for child in tree.find("children"):
            self.import_node(n, child)

    def run(self, node, tree, base=""):
        version = tree.attrib['version']
        xmlbase = tree.attrib['base']
        # import pytest; pytest.set_trace()
        for content in tree.findall("content"):
            self.import_node(node, content)

<commit_msg>redundant whitespace<commit_after>"""
    Import / Export of content

    Requirements:

    - exporting of new types should work, but can be customized
    - xml definitions are versioned, backwards compatibility should be
      present

    - should position be exported? Or should xml order be used, and position
      be renumbered when importing?

    - due to the current nature of nodes, content with a path of /foo/bar may
      exist without a node named /foo. This is an inconsistency in the node tree.
      Such an implementation detail is not relevant to the import/export,
      but how can we handle content without a strict parent? Make it unattached?

    - how to export nodes without content? (unattached nodes)

    - how to handle related content? This can be a User object (its id can
      change) or entirely new models (custom export?)

    Structure (roughly)
    <site version="1" ..>
      <content type=".." slug="..." position="1">
        <created></created>
        <publication></publication>
        <expires></expires>
        <title></title>

        <children>
          <content type=".." slug="..">
          </content>
        </children>
      </content>
      <unattached>

      </unattached>
    </site>

    Should a field be serialized to <title> or <field name="title">? The
    former will not allow for a schema definition
"""

from xml.etree.ElementTree import Element, SubElement, tostring
from django.utils.encoding import smart_unicode
from django.contrib.auth.models import User

from .content import type_registry

class SerializationException(Exception):
    pass

class WheelSerializer(object):
    """ (de) serialize wheel content """
    skip = ('node', )

    def __init__(self):
        pass

    def serialize_owner(self, field, o):
        # import pytest; pytest.set_trace()
        value = getattr(o, field.name, None)
        # value = smart_unicode(value)
        if value is None:
            return ""
        return smart_unicode(value.username)

    def deserialize_owner(self, field, tree):
        username = tree.text
        if not username:
            return None
        return User.objects.get(username=username)

    def serialize(self, spoke):
        o = spoke.instance
        fields = {}
        for field in o._meta.concrete_model._meta.fields:
            handler = getattr(self, "serialize_%s" % field.name, None)
            if field.name in self.skip:
                continue

            if handler:
                fields[field.name] = handler(field, o)
            elif field.serialize:
                if field.rel is None:
                    value = field.value_to_string(o)
                else:
                    value = getattr(o, field.get_attname())
                    value = smart_unicode(value)

                fields[field.name] = value
        xmlfields = Element("fields")
        for k, v in fields.iteritems():
            e = SubElement(xmlfields, "field")
            e.attrib['name'] = k
            e.text = v
        return xmlfields

    def deserialize(self, spoke, tree):
        model = spoke.model
        fields = {}
        # import pytest; pytest.set_trace()
        for field_node in tree.findall("field"):
            field_name = field_node.attrib.get("name")
            if not field_name:
                raise SerializationException("Missing name attribute")
            field = model._meta.get_field(field_name)

            handler = getattr(self, "deserialize_%s" % field.name, None)
            if handler:
                fields[field_name] = handler(field, field_node)

            ##elif field.rel and isinstance(field.rel, models.ManyToManyRel):
            ##    m2m_data[field.name] = self._handle_m2m_field_node(field_node, field)
            ##elif field.rel and isinstance(field.rel, models.ManyToOneRel):
            ##    data[field.attname] = self._handle_fk_field_node(field_node, field)
            else:
                value = field.to_python(field_node.text)
                fields[field.name] = value
        m = model(**fields).save()
        return spoke(m)

class Exporter(object):
    VERSION = 1

    def __init__(self, verbose=0):
        self.verbose = verbose

    def export_node(self, parent, node):
        # import pdb; pdb.set_trace()

        try:
            spoke = node.content().spoke()
        except AttributeError:
            spoke = None

        xmlcontent = SubElement(parent, "content",
                                dict(slug=node.slug(),
                                type=spoke.model.get_name()))
        if spoke:
            contentxml = spoke.serializer().serialize(spoke)
            xmlcontent.append(contentxml)

        children = SubElement(xmlcontent, "children")
        for child in node.children():
            self.export_node(children, child)

    def run(self, node, base="", unattached=True, attachinline=True):
        """ export node and all content beneath it. If unattached is True,
            also export unattached content.

            attachments can be inline or exported to a folder
        """
        root = Element("site")
        root.set('version', str(self.VERSION))
        root.set('base', base)
        self.export_node(root, node)
        return root

class Importer(object):
    def __init__(self, verbose=0):
        self.verbose = verbose

    def import_node(self, node, tree):
        type = tree.attrib['type']
        slug = tree.attrib['slug']
        spoke = type_registry.get(type)
        s = spoke.serializer().deserialize(spoke, tree)
        if slug == "":
            n = node
        else:
            n = node.add(slug)
        n.set(s.instance)

        for child in tree.find("children"):
            self.import_node(n, child)

    def run(self, node, tree, base=""):
        version = tree.attrib['version']
        xmlbase = tree.attrib['base']
        # import pytest; pytest.set_trace()
        for content in tree.findall("content"):
            self.import_node(node, content)

<|endoftext|>


<commit_before>from pupa.scrape import Scraper, Legislator
from pupa.models import Organization
from utils import lxmlize

import re
import urllib2
import os
import subprocess

COUNCIL_PAGE = 'http://www.community.gov.yk.ca/pdf/loc_govdir.pdf'

class YukonMunicipalitiesPersonScraper(Scraper):

  def get_people(self):
    response = urllib2.urlopen(COUNCIL_PAGE).read()
    pdf = open('/tmp/yt.pdf', 'w')
    pdf.write(response)
    pdf.close()

    data = subprocess.check_output(['pdftotext', '-layout', '/tmp/yt.pdf', '-'])
    data = re.split(r'\n\s*\n', data)
    for municipality in data:

      if not 'Councillors' in municipality:
        continue
      lines = municipality.split('\n')
      if 'Page' in lines[0]:
        lines.pop(0)
        if not lines[0].strip():
          lines.pop(0)
      col1end = re.search(r'\s{2,}(\w)', lines[0].strip()).end()
      col2end = re.search(r':\s{2,}(\w)', lines[0].strip()).end()

      if 'Council' in lines[1]:
        address = lines[2][:col1end - 1].strip() + ' ' + lines[3][:col1end - 1].strip()
        district = lines[0][:col1end - 1].strip() + ' ' + lines[1][:col1end - 1].strip()
      else:
        address = lines[1][:col1end - 1].strip() + ' ' + lines[2][:col1end - 1].strip()
        district = lines[0][:col1end - 1].strip()

      chamber = district + ' Council'
      organization = Organization(name=chamber, chamber=chamber, classification='legislature', jurisdiction_id=self.jurisdiction.jurisdiction_id)
      organization.add_source(COUNCIL_PAGE)
      yield organization

      phone = re.findall(r'(?<=Phone: )\(?(\d{3}[\)-] ?\d{3}-\d{4})', municipality)[0].replace(') ', '-')
      email = re.findall(r'(?<=E-mail:) (\S*)', municipality)[0]
      fax = None
      if 'Fax' in municipality:
        fax = re.findall(r'(?<=Fax: )\(?(\d{3}[\)-] ?\d{3}-\d{4})', municipality)[0].replace(') ', '-')
      website = None
      if 'Website' in municipality:
        website = re.findall(r'((http:\/\/|www.)(\S*))', municipality)[0][0]

      councillor_or_mayor = False
      for line in lines:
        if 'Mayor:' in line:
          councillor_or_mayor = True
          role = 'mayor'
          continue
        if 'Councillors' in line:
          councillor_or_mayor = True
          role = 'councillor'
          continue
        if councillor_or_mayor:
          councillor = line[col1end - 1:col2end - 1].strip()
          if not councillor:
            continue
          p = Legislator(name=councillor, post_id=district)
          p.add_source(COUNCIL_PAGE)
          p.add_membership(organization, role=role, chamber=chamber)
          p.add_contact('address', address, None)
          p.add_contact('voice', phone, None)
          p.add_contact('email', email, None)
          if fax:
            p.add_contact('fax', fax, None)
          if website:
            p.add_link(website, None)
          yield p

    os.system('rm yt.pdf')
<commit_msg>fix yt contact_details notes<commit_after>from pupa.scrape import Scraper, Legislator
from pupa.models import Organization
from utils import lxmlize

import re
import urllib2
import os
import subprocess

COUNCIL_PAGE = 'http://www.community.gov.yk.ca/pdf/loc_govdir.pdf'

class YukonMunicipalitiesPersonScraper(Scraper):

  def get_people(self):
    response = urllib2.urlopen(COUNCIL_PAGE).read()
    pdf = open('/tmp/yt.pdf', 'w')
    pdf.write(response)
    pdf.close()

    data = subprocess.check_output(['pdftotext', '-layout', '/tmp/yt.pdf', '-'])
    data = re.split(r'\n\s*\n', data)
    for municipality in data:

      if not 'Councillors' in municipality:
        continue
      lines = municipality.split('\n')
      if 'Page' in lines[0]:
        lines.pop(0)
        if not lines[0].strip():
          lines.pop(0)
      col1end = re.search(r'\s{2,}(\w)', lines[0].strip()).end()
      col2end = re.search(r':\s{2,}(\w)', lines[0].strip()).end()

      if 'Council' in lines[1]:
        address = lines[2][:col1end - 1].strip() + ' ' + lines[3][:col1end - 1].strip()
        district = lines[0][:col1end - 1].strip() + ' ' + lines[1][:col1end - 1].strip()
      else:
        address = lines[1][:col1end - 1].strip() + ' ' + lines[2][:col1end - 1].strip()
        district = lines[0][:col1end - 1].strip()

      chamber = district + ' Council'
      organization = Organization(name=chamber, chamber=chamber, classification='legislature', jurisdiction_id=self.jurisdiction.jurisdiction_id)
      organization.add_source(COUNCIL_PAGE)
      yield organization

      phone = re.findall(r'(?<=Phone: )\(?(\d{3}[\)-] ?\d{3}-\d{4})', municipality)[0].replace(') ', '-')
      email = re.findall(r'(?<=E-mail:) (\S*)', municipality)[0]
      fax = None
      if 'Fax' in municipality:
        fax = re.findall(r'(?<=Fax: )\(?(\d{3}[\)-] ?\d{3}-\d{4})', municipality)[0].replace(') ', '-')
      website = None
      if 'Website' in municipality:
        website = re.findall(r'((http:\/\/|www.)(\S*))', municipality)[0][0]

      councillor_or_mayor = False
      for line in lines:
        if 'Mayor:' in line:
          councillor_or_mayor = True
          role = 'mayor'
          continue
        if 'Councillors' in line:
          councillor_or_mayor = True
          role = 'councillor'
          continue
        if councillor_or_mayor:
          councillor = line[col1end - 1:col2end - 1].strip()
          if not councillor:
            continue
          p = Legislator(name=councillor, post_id=district)
          p.add_source(COUNCIL_PAGE)
          p.add_membership(organization, role=role, chamber=chamber)
          p.add_contact('address', address, 'legislature')
          p.add_contact('voice', phone, 'legislature')
          p.add_contact('email', email, None)
          if fax:
            p.add_contact('fax', fax, 'legislature')
          if website:
            p.add_link(website, None)
          yield p

    os.system('rm yt.pdf')
<|endoftext|>


<commit_before>import mimetypes
import re
import os
import datetime

from django.db import models, IntegrityError
from django.utils import timezone
from django.contrib.auth.models import User, Group

from userena.models import UserenaLanguageBaseProfile

from drole.fields import RoleField

from .impexp import WheelSerializer
from .themes import theme_registry

## import pytz

from django.utils.translation import ugettext as _

class WheelProfile(UserenaLanguageBaseProfile):
    """ timezone, ... ?
    """
    ##timezone = models.CharField(max_length=100, default=config.DEFAULT_TIMEZONE,
    ##    choices=[(x, x) for x in pytz.common_timezones])
    inform = models.BooleanField(default=False)
    user = models.OneToOneField(User,
                                unique=True,
                                verbose_name=_('user'),
                                related_name='my_profile')

    twitter = models.TextField(blank=True, null=False, default="")
    google = models.TextField(blank=True, null=False, default="")
    linkedin = models.TextField(blank=True, null=False, default="")

## backwards compat, should import from node directly
from .node import NodeException, DuplicatePathException, InvalidPathException
from .node import NodeInUse, CantRenameRoot, NodeBase
from .node import WHEEL_NODE_BASECLASS
from .node import Node

from .content import ContentClass, ContentBase
from .content import WHEEL_CONTENT_BASECLASS
from .content import Content, ClassContentManager
from .content import FileContent, ImageContent
from .content import type_registry

class Configuration(models.Model):
    title = models.CharField(max_length=256, blank=True, null=False, default="")
    description = models.TextField(blank=True, null=False, default="")
    theme = models.CharField(max_length=256, blank=True, null=False, default="default")

    analytics = models.CharField(max_length=50, blank=True, null=False, default="")
    head = models.TextField(blank=True, null=False, default="")

    sender = models.CharField(max_length=100, blank=True, null=False, default="")
    sendermail = models.EmailField(max_length=100, blank=True, null=False, default="")
    mailto = models.EmailField(max_length=100, blank=True, null=False, default="")

    @classmethod
    def config(cls):
        """ singleton-ish pattern """
        try:
            instance = Configuration.objects.all()[0]
        except IndexError:
            instance = Configuration()
            instance.save()
        return instance

    def themeinfo(self):
        """ resolve self.theme into a Theme instance """
        return theme_registry.find(self.theme)

class Role(models.Model):
    """
        System-wide (non-local) roles for users
    """
    role = RoleField(max_length=255, blank=False)
    user = models.ForeignKey(User, related_name="roles", null=True, blank=True)
    group = models.ForeignKey(Group, related_name="roles", null=True, blank=True)

#class ConfigItem(models.Model):
#    name = models.CharField(max_length=256, blank=False, null=False)
#    value = models.TextField(blank=True, null=False, default="")
#    type = models.CharField(max_length=256, blank=False, null=False,
#                            default="string")
#    ns = models.CharField(max_length=256, blank=True, null=False,
#                          default="default")
#    configuration = models.ForeignKey(Configuration, related_name="items",
#                           blank=False, null=False)

## signals for login/logout logging

from django.db.models.signals import post_save
from django.dispatch import receiver
from userena.signals import signup_complete
from django.contrib.auth.signals import user_logged_in, user_logged_out

import stracks
## django 1.5
## from django.contrib.auth.signals import user_login_failed
##
## @receiver(user_login_failed)
## def log_failure(sender, credentials, **kwargs):
##     """ Log the user logging out """
##     stracksclient.warning("Authentication failed", data=credentials)
##

from django.db.utils import DatabaseError
import logging

@receiver(post_save, sender=User)
def log_create(sender, instance, created, **kwargs):
    """ Log the creation of a new user """
    if created:
        stracks.user(instance).log("? has been created")
    ## make sure it has a wheel profile
    try:
        WheelProfile.objects.get_or_create(user=instance)
    except DatabaseError:
        logging.error("Failed to create profile for %s, perhaps migrations haven't run yet?" % instance)
        from django.db import connection
        connection._rollback()

from south.signals import post_migrate

@receiver(post_migrate)
def create_profiles(app, **kwargs):
    if app == "wheelcms_axle":
        for u in User.objects.all():
            WheelProfile.objects.get_or_create(user=u)

@receiver(signup_complete, dispatch_uid='stracks.log_signup')
def log_signup(sender, signal, user, **kwargs):
    """ Log the creation of a new user """
    stracks.user(user).log("? has completed (userena) signup")

@receiver(user_logged_in, dispatch_uid='stracks.log_signin')
def log_login(sender, request, user, **kwargs):
    """ Log the user logging in """
    stracks.user(user).log("? has logged in", action=stracks.login())

@receiver(user_logged_out, dispatch_uid='stracks.log_signout')
def log_logout(sender, request, user, **kwargs):
    """ Log the user logging out """
    stracks.user(user).log("? has logged out", action=stracks.logout())

@receiver(post_save, dispatch_uid="wheelcms_axle.spoke.assign_perms")
def assign_perms(sender, instance, created, **kwargs):

    if issubclass(sender, Content):
        spoke = instance.spoke()
        if spoke:
            spoke.assign_perms()
            if created:
                wf = spoke.workflow()
                state = instance.state
                assignment = wf.permission_assignment.get(state)
                if assignment:
                    spoke.update_perms(assignment)

    if hasattr(sender, 'permission_assignment') and created:
        assign_perms(instance, instance.permission_assignment)

<commit_msg>make sure a default workflow is set<commit_after>import mimetypes
import re
import os
import datetime

from django.db import models, IntegrityError
from django.utils import timezone
from django.contrib.auth.models import User, Group

from userena.models import UserenaLanguageBaseProfile

from drole.fields import RoleField

from .impexp import WheelSerializer
from .themes import theme_registry

## import the default workflow, so it can register itself as default
## should be in some sort of base configuration class in stead of here

import workflows.default

## import pytz

from django.utils.translation import ugettext as _

class WheelProfile(UserenaLanguageBaseProfile):
    """ timezone, ... ?
    """
    ##timezone = models.CharField(max_length=100, default=config.DEFAULT_TIMEZONE,
    ##    choices=[(x, x) for x in pytz.common_timezones])
    inform = models.BooleanField(default=False)
    user = models.OneToOneField(User,
                                unique=True,
                                verbose_name=_('user'),
                                related_name='my_profile')

    twitter = models.TextField(blank=True, null=False, default="")
    google = models.TextField(blank=True, null=False, default="")
    linkedin = models.TextField(blank=True, null=False, default="")

## backwards compat, should import from node directly
from .node import NodeException, DuplicatePathException, InvalidPathException
from .node import NodeInUse, CantRenameRoot, NodeBase
from .node import WHEEL_NODE_BASECLASS
from .node import Node

from .content import ContentClass, ContentBase
from .content import WHEEL_CONTENT_BASECLASS
from .content import Content, ClassContentManager
from .content import FileContent, ImageContent
from .content import type_registry

class Configuration(models.Model):
    title = models.CharField(max_length=256, blank=True, null=False, default="")
    description = models.TextField(blank=True, null=False, default="")
    theme = models.CharField(max_length=256, blank=True, null=False, default="default")

    analytics = models.CharField(max_length=50, blank=True, null=False, default="")
    head = models.TextField(blank=True, null=False, default="")

    sender = models.CharField(max_length=100, blank=True, null=False, default="")
    sendermail = models.EmailField(max_length=100, blank=True, null=False, default="")
    mailto = models.EmailField(max_length=100, blank=True, null=False, default="")

    @classmethod
    def config(cls):
        """ singleton-ish pattern """
        try:
            instance = Configuration.objects.all()[0]
        except IndexError:
            instance = Configuration()
            instance.save()
        return instance

    def themeinfo(self):
        """ resolve self.theme into a Theme instance """
        return theme_registry.find(self.theme)

class Role(models.Model):
    """
        System-wide (non-local) roles for users
    """
    role = RoleField(max_length=255, blank=False)
    user = models.ForeignKey(User, related_name="roles", null=True, blank=True)
    group = models.ForeignKey(Group, related_name="roles", null=True, blank=True)

#class ConfigItem(models.Model):
#    name = models.CharField(max_length=256, blank=False, null=False)
#    value = models.TextField(blank=True, null=False, default="")
#    type = models.CharField(max_length=256, blank=False, null=False,
#                            default="string")
#    ns = models.CharField(max_length=256, blank=True, null=False,
#                          default="default")
#    configuration = models.ForeignKey(Configuration, related_name="items",
#                           blank=False, null=False)

## signals for login/logout logging

from django.db.models.signals import post_save
from django.dispatch import receiver
from userena.signals import signup_complete
from django.contrib.auth.signals import user_logged_in, user_logged_out

import stracks
## django 1.5
## from django.contrib.auth.signals import user_login_failed
##
## @receiver(user_login_failed)
## def log_failure(sender, credentials, **kwargs):
##     """ Log the user logging out """
##     stracksclient.warning("Authentication failed", data=credentials)
##

from django.db.utils import DatabaseError
import logging

@receiver(post_save, sender=User)
def log_create(sender, instance, created, **kwargs):
    """ Log the creation of a new user """
    if created:
        stracks.user(instance).log("? has been created")
    ## make sure it has a wheel profile
    try:
        WheelProfile.objects.get_or_create(user=instance)
    except DatabaseError:
        logging.error("Failed to create profile for %s, perhaps migrations haven't run yet?" % instance)
        from django.db import connection
        connection._rollback()

from south.signals import post_migrate

@receiver(post_migrate)
def create_profiles(app, **kwargs):
    if app == "wheelcms_axle":
        for u in User.objects.all():
            WheelProfile.objects.get_or_create(user=u)

@receiver(signup_complete, dispatch_uid='stracks.log_signup')
def log_signup(sender, signal, user, **kwargs):
    """ Log the creation of a new user """
    stracks.user(user).log("? has completed (userena) signup")

@receiver(user_logged_in, dispatch_uid='stracks.log_signin')
def log_login(sender, request, user, **kwargs):
    """ Log the user logging in """
    stracks.user(user).log("? has logged in", action=stracks.login())

@receiver(user_logged_out, dispatch_uid='stracks.log_signout')
def log_logout(sender, request, user, **kwargs):
    """ Log the user logging out """
    stracks.user(user).log("? has logged out", action=stracks.logout())

@receiver(post_save, dispatch_uid="wheelcms_axle.spoke.assign_perms")
def assign_perms(sender, instance, created, **kwargs):

    if issubclass(sender, Content):
        spoke = instance.spoke()
        if spoke:
            spoke.assign_perms()
            if created:
                wf = spoke.workflow()
                state = instance.state
                assignment = wf.permission_assignment.get(state)
                if assignment:
                    spoke.update_perms(assignment)

    if hasattr(sender, 'permission_assignment') and created:
        assign_perms(instance, instance.permission_assignment)

<|endoftext|>


<commit_before>"""Load test data"""
import math
import heap

def shouldBalance(left, right):
    """Check if 2 heaps are more than 1 node apart"""
    return math.fabs(left.size() - right.size()) > 1

def balance(left, right):
    """Balance two heaps that are off by 1 value"""
    if left.size() > right.size():
        temp = left.remove()
        right.insert(temp)
    elif left.size() < right.size():
        temp = right.remove()
        left.insert(temp)
    else:
        raise ValueError("Heaps were of same size")

    if shouldBalance(left, right):
        raise ValueError("Balances was called too late")

    return True

FILENAME = "test.txt"
# file = "median.txt"

data = []
with open(FILENAME) as f:
    for num in f:
        data.append(int(num))

median = None
leftHeap = heap.Heap(heap.MAX)
rightHeap = heap.Heap(heap.MIN)

for i, d in enumerate(data):
    if d < median:
        leftHeap.insert(d)
    else:
        rightHeap.insert(d)
    # print "Left"
    # leftHeap.show()
    # print "Right"
    # rightHeap.show()

    if shouldBalance(leftHeap, rightHeap):
        balance(leftHeap, rightHeap)

    if leftHeap.size() > rightHeap.size():
        median = leftHeap.top()
    elif rightHeap.size() > leftHeap.size():
        median = rightHeap.top()
    else:
        median = (leftHeap.top() + rightHeap.top()) / 2.0

    print float(median)
<commit_msg>Works for large data set<commit_after>"""Load test data"""
import math
import heap

def shouldBalance(left, right):
    """Check if 2 heaps are more than 1 node apart"""
    return math.fabs(left.size() - right.size()) > 1

def balance(left, right):
    """Balance two heaps that are off by 1 value"""
    if left.size() > right.size():
        temp = left.remove()
        right.insert(temp)
    elif left.size() < right.size():
        temp = right.remove()
        left.insert(temp)
    else:
        raise ValueError("Heaps were of same size")

    if shouldBalance(left, right):
        raise ValueError("Balances was called too late")

    return True

# FILENAME = "test.txt"
FILENAME = "median.txt"

data = []
with open(FILENAME) as f:
    for num in f:
        data.append(int(num))

median = None
leftHeap = heap.Heap(heap.MAX)
rightHeap = heap.Heap(heap.MIN)
medianSum = 0
for i, d in enumerate(data):
    if d < median:
        leftHeap.insert(d)
    else:
        rightHeap.insert(d)

    if shouldBalance(leftHeap, rightHeap):
        balance(leftHeap, rightHeap)

    if leftHeap.size() > rightHeap.size():
        median = leftHeap.top()
    elif rightHeap.size() > leftHeap.size():
        median = rightHeap.top()
    else:
        median = leftHeap.top()
    medianSum += median

print medianSum
print medianSum % 10000
<|endoftext|>