gc_area/create_polygon.py

#!/bin/python3

import argparse
import osmapi
import shapely
import pickle


def chunks(lst, n):
    """Yield successive n-sized chunks from lst.
    https://stackoverflow.com/a/312464
    CC BY-SA 4.0
    """
    for i in range(0, len(lst), n):
        yield lst[i:i + n]


def get_shell_ids(osm_api, ways):
    """
    Get a list of OSM node IDs
    that forms the shell of the area's polygon.
    """
    shell_ids = list()
    for way in ways:
        way = osm_api.WayGet(way['ref'])
        node_ids = way['nd']
        if not shell_ids:
            shell_ids.append(node_ids[0])
        for node_id in node_ids[1:]:
            shell_ids.append(node_id)

    return shell_ids


def get_node_data(osm_api, shell_ids):
    """Get the data of OSM nodes as a dict of (id -> data)."""
    node_data = dict()
    for chunk in chunks(shell_ids, 100):
        node_data |= osm_api.NodesGet(chunk)

    return node_data


def get_shell(shell_ids, node_data):
    """Get the shell of the area's polygon as a list of (lat, lon) tuples."""
    shell = list()
    for shell_id in shell_ids:
        node = node_data[shell_id]
        shell.append((node['lat'], node['lon']))

    return shell


def main(args):
    osm_api = osmapi.OsmApi()
    rel = osm_api.RelationGet(args.osm_relation_id)
    # Filter out city centre etc.
    ways = filter(lambda x: x['type'] == 'way', rel['member'])
    shell_ids = get_shell_ids(osm_api, ways)
    node_data = get_node_data(osm_api, shell_ids)
    shell = get_shell(shell_ids, node_data)
    polygon = shapely.Polygon(shell)

    with open(args.output_file, 'wb') as f:
        pickle.dump(polygon, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="""
    Creates a polygon from an OpenStreetMap relation and saves this
    to a file as a Python pickle.
    """)
    parser.add_argument("osm-relation-id",
                        help="OpenStreetMap relation ID", type=int)
    parser.add_argument(
        "-v", "--output-file", help="Output file for the resulting polygon",
        default="polygon.pickle")
    args = parser.parse_args()

    main(args)