Spatial Queries in Python

spatialquery.py

import fiona
import os
import rtree

from collections import defaultdict
from shapely.geometry import shape

CONTAINS = 'contains'
CROSSES = 'crosses'
DISJOINT = 'disjoint'
EQUALS = 'equals'
INTERSECTS = 'intersects'
TOUCHES = 'touches'
WITHIN = 'within'


def generate_index(records, index_path=None):
    prop = rtree.index.Property()
    if index_path is not None:
        prop.storage = rtree.index.RT_Disk
        prop.overwrite = index_path

    sp_index = rtree.index.Index(index_path, properties=prop)
    for n, ft in enumerate(records):
        if ft['geometry'] is not None:
            sp_index.insert(n, shape(ft['geometry']).bounds)
    return sp_index


def query(layer1, layer2, spatial_predicate=INTERSECTS, group_field=None):
    result = defaultdict(list)
    with fiona.open(layer2) as datasrc2:
        fname, _ = os.path.splitext(layer2)
        if os.path.exists(fname + '.dat') and os.path.exists(fname + '.idx'):
            layer2_index = rtree.index.Index(fname)
        else:
            layer2_index = generate_index(datasrc2, fname)

        with fiona.open(layer1) as datasrc1:
            for ft1 in datasrc1:
                if ft1['geometry'] is None:
                    continue
                shp1 = shape(ft1['geometry'])
                for fid in layer2_index.intersection(shp1.bounds):
                    shp2 = shape(datasrc2[fid]['geometry'])
                    if getattr(shp1, spatial_predicate)(shp2):
                        if group_field:
                            result[ft1['properties'][group_field]].append(fid)
                        else:
                            result[ft1['id']].append(fid)
    return result