<commit_before>import fiona
import os
import rtree
import ujson as json

from collections import OrderedDict
from shapely.geometry import Point, Polygon, MultiPolygon
from shapely.prepared import prep
from shapely.geometry.geo import mapping

DEFAULT_INDEX_FILENAME = 'rtree'
DEFAULT_POLYS_FILENAME = 'polygons.geojson'


class RTreePolygonIndex(object):
    include_only_properties = None

    def __init__(self, index=None, polygons=None, save_dir=None,
                 index_filename=None,
                 include_only_properties=None):
        if save_dir:
            self.save_dir = save_dir
        else:
            self.save_dir = None

        if index_filename:
            self.index_path = os.path.join(save_dir or '.', index_filename)
        else:
            self.index_path = None

        if not index and self.index_path:
            self.index = rtree.index.Index(self.index_path)
        elif not index:
            self.index = rtree.index.Index()
        else:
            self.index = index

        if include_only_properties and hasattr(include_only_properties, '__contains__'):
            self.include_only_properties = include_only_properties

        if not polygons:
            self.polygons = []
        else:
            self.polygons = polygons

    def index_polygon(self, id, polygon):
        self.index.insert(id, polygon.bounds)

    def add_polygon(self, poly, properties, simplify_tolerance=0.0001, preserve_topology=True):
        poly = poly.simplify(simplify_tolerance, preserve_topology=preserve_topology)
        if self.include_only_properties:
            properties = {k: v for k, v in properties.iteritems() if k in self.include_only_properties}

        self.polygons.append((properties, prep(poly)))

    @classmethod
    def create_from_shapefiles(cls, inputs, output_dir,
                               index_filename=DEFAULT_INDEX_FILENAME,
                               polys_filename=DEFAULT_POLYS_FILENAME):
        index = cls(save_dir=output_dir, index_filename=index_filename)
        for input_file in inputs:
            f = fiona.open(input_file)

            i = 0

            for rec in f:
                if not rec or not rec.get('geometry') or 'type' not in rec['geometry']:
                    continue
                poly_type = rec['geometry']['type']
                if poly_type == 'Polygon':
                    poly = Polygon(rec['geometry']['coordinates'][0])
                    index.index_polygon(i, poly)
                    index.add_polygon(poly, rec['properties'])
                elif poly_type == 'MultiPolygon':
                    polys = []
                    for coords in rec['geometry']['coordinates']:
                        poly = Polygon(coords[0])
                        polys.append(poly)
                        index.index_polygon(i, poly)

                    index.add_polygon(MultiPolygon(polys), rec['properties'])
                i += 1
        return index

    def save_polygons(self, out_filename):
        out = open(out_filename, 'w')
        features = []
        for props, poly in self.polygons:
            features.append({
                'type': 'Feature',
                'geometry': mapping(poly.context),
                'properties': props
            })
        json.dump({'type': 'FeatureCollection',
                   'features': features},
                  out)

    def save(self, polys_filename=DEFAULT_POLYS_FILENAME):
        self.save_polygons(os.path.join(self.save_dir, polys_filename))
        # need to close index before loading it
        self.index.close()

    @classmethod
    def load_polygons(cls, filename):
        feature_collection = json.load(open(filename))
        polygons = []
        for feature in feature_collection['features']:
            poly_type = feature['geometry']['type']

            if poly_type == 'Polygon':
                poly = Polygon(feature['geometry']['coordinates'][0])
                polygons.append((feature['properties'], prep(poly)))
            elif poly_type == 'MultiPolygon':
                polys = []
                for coords in feature['geometry']['coordinates']:
                    poly = Polygon(coords[0])
                    polys.append(poly)

                polygons.append((feature['properties'], prep(MultiPolygon(polys))))
        return polygons

    @classmethod
    def load(cls, d, index_name=DEFAULT_INDEX_FILENAME, polys_filename=DEFAULT_POLYS_FILENAME):
        index = rtree.index.Index(os.path.join(d, index_name))
        polys = cls.load_polygons(os.path.join(d, polys_filename))
        return cls(index=index, polygons=polys, save_dir=d)

    def get_candidate_polygons(self, lat, lon):
        return OrderedDict.fromkeys(self.index.intersection((lon, lat, lon, lat))).keys()

    def point_in_poly(self, lat, lon):
        polys = self.get_candidate_polygons(lat, lon)
        if len(polys) == 1:
            return self.polygons[polys[0]][0]
        pt = Point(lon, lat)
        for i in polys:
            props, poly = self.polygons[i]
            if poly.contains(pt):
                return props<commit_msg>[fix] RTreePolygonIndex should be doing an exact point-in-polygon test even if there are no collisions since R-trees only really store a bounding box.<commit_after>import fiona
import os
import rtree
import ujson as json

from collections import OrderedDict
from shapely.geometry import Point, Polygon, MultiPolygon
from shapely.prepared import prep
from shapely.geometry.geo import mapping

DEFAULT_INDEX_FILENAME = 'rtree'
DEFAULT_POLYS_FILENAME = 'polygons.geojson'


class RTreePolygonIndex(object):
    include_only_properties = None

    def __init__(self, index=None, polygons=None, save_dir=None,
                 index_filename=None,
                 include_only_properties=None):
        if save_dir:
            self.save_dir = save_dir
        else:
            self.save_dir = None

        if index_filename:
            self.index_path = os.path.join(save_dir or '.', index_filename)
        else:
            self.index_path = None

        if not index and self.index_path:
            self.index = rtree.index.Index(self.index_path)
        elif not index:
            self.index = rtree.index.Index()
        else:
            self.index = index

        if include_only_properties and hasattr(include_only_properties, '__contains__'):
            self.include_only_properties = include_only_properties

        if not polygons:
            self.polygons = []
        else:
            self.polygons = polygons

    def index_polygon(self, id, polygon):
        self.index.insert(id, polygon.bounds)

    def add_polygon(self, poly, properties, simplify_tolerance=0.0001, preserve_topology=True):
        poly = poly.simplify(simplify_tolerance, preserve_topology=preserve_topology)
        if self.include_only_properties:
            properties = {k: v for k, v in properties.iteritems() if k in self.include_only_properties}

        self.polygons.append((properties, prep(poly)))

    @classmethod
    def create_from_shapefiles(cls, inputs, output_dir,
                               index_filename=DEFAULT_INDEX_FILENAME,
                               polys_filename=DEFAULT_POLYS_FILENAME):
        index = cls(save_dir=output_dir, index_filename=index_filename)
        for input_file in inputs:
            f = fiona.open(input_file)

            i = 0

            for rec in f:
                if not rec or not rec.get('geometry') or 'type' not in rec['geometry']:
                    continue
                poly_type = rec['geometry']['type']
                if poly_type == 'Polygon':
                    poly = Polygon(rec['geometry']['coordinates'][0])
                    index.index_polygon(i, poly)
                    index.add_polygon(poly, rec['properties'])
                elif poly_type == 'MultiPolygon':
                    polys = []
                    for coords in rec['geometry']['coordinates']:
                        poly = Polygon(coords[0])
                        polys.append(poly)
                        index.index_polygon(i, poly)

                    index.add_polygon(MultiPolygon(polys), rec['properties'])
                i += 1
        return index

    def save_polygons(self, out_filename):
        out = open(out_filename, 'w')
        features = []
        for props, poly in self.polygons:
            features.append({
                'type': 'Feature',
                'geometry': mapping(poly.context),
                'properties': props
            })
        json.dump({'type': 'FeatureCollection',
                   'features': features},
                  out)

    def save(self, polys_filename=DEFAULT_POLYS_FILENAME):
        self.save_polygons(os.path.join(self.save_dir, polys_filename))
        # need to close index before loading it
        self.index.close()

    @classmethod
    def load_polygons(cls, filename):
        feature_collection = json.load(open(filename))
        polygons = []
        for feature in feature_collection['features']:
            poly_type = feature['geometry']['type']

            if poly_type == 'Polygon':
                poly = Polygon(feature['geometry']['coordinates'][0])
                polygons.append((feature['properties'], prep(poly)))
            elif poly_type == 'MultiPolygon':
                polys = []
                for coords in feature['geometry']['coordinates']:
                    poly = Polygon(coords[0])
                    polys.append(poly)

                polygons.append((feature['properties'], prep(MultiPolygon(polys))))
        return polygons

    @classmethod
    def load(cls, d, index_name=DEFAULT_INDEX_FILENAME, polys_filename=DEFAULT_POLYS_FILENAME):
        index = rtree.index.Index(os.path.join(d, index_name))
        polys = cls.load_polygons(os.path.join(d, polys_filename))
        return cls(index=index, polygons=polys, save_dir=d)

    def get_candidate_polygons(self, lat, lon):
        return OrderedDict.fromkeys(self.index.intersection((lon, lat, lon, lat))).keys()

    def point_in_poly(self, lat, lon):
        polys = self.get_candidate_polygons(lat, lon)
        pt = Point(lon, lat)
        for i in polys:
            props, poly = self.polygons[i]
            if poly.contains(pt):
                return props
        return None
{'fraction_non_alphanumeric': 0.09771169257507, 'fraction_numerical': 0.002413826397605484, 'mean_word_length': 2.101197604790419, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5882734', 'n_tokens_mistral': 3010, 'n_tokens_neox': 2840, 'n_words': 630}