Sorted Countries Demo for FOSS4G

sorted.py

#!/usr/bin/python

import shapely, json, math, time
from shapely import speedups, wkb
from shapely.geometry import LineString, Point, Polygon
from osgeo import ogr
from optparse import OptionParser
from math import atan2, degrees, sin, cos

# Be fast if possible.
if shapely.speedups.available:
    shapely.speedups.enable()

parser = OptionParser(usage='%prog output.shp input.shp')

# Plug Shapely and OGR into each other and hope it works.
def load_shapefile(filename):
    # Abstracting these calls into another function
    # will segfault due to Python/OGR fights
    try:
        source = ogr.Open(filename)
        l = source.GetLayer(0)
    except:
        raise Exception('Error opening input file')

    f = l.GetNextFeature()
    if not f:
        raise Exception('empty shapefile found')
    f = f.Clone()
    while f:
        # yielding here in an early attempt to avoid
        # putting everything in memory. But since we're dealing
        # with min/max measures, that's not a possibility, yet.
        yield [f.GetField('ISO_A3'),
                f.GetField('NAME'),
                shapely.wkb.loads(f.GetGeometryRef().ExportToWkb())]
        f = l.GetNextFeature()
        f = f.Clone() if f else False

def zeroed(feat):
    iso3, name, p = feat
    x, y = p.exterior.xy
    mx = min(x)
    my = min(y)
    nx = map(lambda j: j - mx, x)
    ny = map(lambda j: j - my, y)
    return [iso3, name, Polygon(zip(nx, ny))]

def pushto(feat, dx, dy):
    iso3, name, p = feat
    x, y = p.exterior.xy
    nx = map(lambda j: j + dx, x)
    ny = map(lambda j: j + dy, y)
    return [iso3, name, Polygon(zip(nx, ny))]

if __name__ == "__main__":
    (options, args) = parser.parse_args()
    (outfile, infile) = args

    shpdriver = ogr.GetDriverByName('ESRI Shapefile')

    features = sorted([x for x in load_shapefile(infile)], key=lambda a: a[2].area)

    zeros = map(zeroed, filter(lambda x: x[2].geometryType() == 'Polygon', features))
    pushed = []

    squiggle_ds = shpdriver.CreateDataSource(outfile)
    lyr = squiggle_ds.CreateLayer('sorted')

    field_defn = ogr.FieldDefn('ISO3', ogr.OFTString)
    lyr.CreateField(field_defn)

    field_defn = ogr.FieldDefn('NAME', ogr.OFTString)
    lyr.CreateField(field_defn)

    positions = []

    dx = 0
    dy = 0
    for j in zeros:
        jsonpos = [-175, -70]
        j = pushto(j, -175, -70)
        if len(pushed):

            if (dx > 280):
                dx = 0
                dy = dy + 10

            j = pushto(j, dx, dy)
            jsonpos = [jsonpos[0] + dx, jsonpos[1] + dy]
            while pushed[-1][2].intersects(j[2]) or (len(pushed) > 1 and pushed[-2][2].intersects(j[2])) and dx < 200:
                dx = dx + 20
                j = pushto(j, 20, 0)
                jsonpos = [jsonpos[0] + 20, jsonpos[1] + 0]

        pushed.append(j)
        positions.append({
            "iso3": pushed[-1][0],
            "name": pushed[-1][1],
            "position": jsonpos
            })
        q_geom = ogr.CreateGeometryFromWkb(pushed[-1][2].wkb)
        q_feature = ogr.Feature(lyr.GetLayerDefn())
        q_feature.SetField(0, pushed[-1][0])
        q_feature.SetField(1, pushed[-1][1])
        q_feature.SetGeometry(q_geom)
        lyr.CreateFeature(q_feature)
    json.dump(positions, open('positions.json', 'w'))