class TrigramGeocoder(object): def __init__(self, name, srid=4326, n=3, fc=None, clear=False): self.code_store = GeoJSONCollection( db=settings.MONGODB_ROUTES['ga_geocoder'] if 'ga_geocoder' in settings.MONGODB_ROUTES else settings.MONGODB_ROUTES['default'], collection=name, index_path=app_settings.GEO_INDEX_PATH, srid=srid, fc=fc, clear=clear ) self.ngram_store = settings.MONGODB_ROUTES['ga_geocoder'][name + "_ngrams"] self.code_store['kind'] = 'trigram' self.code_store['ngrams'] = name + "_ngrams" self.serialize = lambda x: x.json self.deserialize = lambda x: GEOSGeometry(x, srid=self.code_store.srid) self.parser = parsers.en_us.address_trigrams def __getitem__(self, code): """ Returns a feature if an exact match is found or a feature collection if an approximate match is found :param code: :return: """ srid=None if isinstance(code, tuple): code, srid = code val = self.code_store.coll.find_one(code) if val: if srid: g = GEOSGeometry(json.dumps(val['geometry']), srid=self.code_store.srid) g.transform(srid) val['geometry'] = json.loads(g.json) return val else: val = { 'type' : "FeatureCollection" } features = [] ngrams = self.parser(code) scores = {} for ngram in ngrams: counts = self.ngram_store.find({'ngram' : ngram}) for count in counts: if count['code'] not in scores: scores[ count['code'] ] += count['count'] for code, ct in sorted(scores.items(), key=lambda x,y: y): features.append(self.code_store.coll.find_one(code)) if len(features): val['features'] = features return val else: raise KeyError(code) def __delitem__(self, code): ngrams = self.parser(code) if self.code_store.coll.find_one(code): self.code_store.coll.remove(code) for ngram in ngrams: doc = self.ngram_store.find_one(ngram) doc['count'] -= 1 self.ngram_store.save(doc) def __setitem__(self, code, geom, geom_serializer=None): del self[code] fc = { "type" : "GeometryCollection" } features = [] index = {} ngrams = self.parser(code) for ngram in ngrams: if not ngram in index: index[ngram] = Counter() index[ngram][code] += 1 if geom_serializer: geom = json.loads(geom_serializer(geom)) else: geom = json.loads(self.serialize(geom)) orig = code features.append({ '_id' : code, 'geometry' : geom, 'properties' : { 'name' : orig } }) self.code_store.insert_features(fc) for ngram, counter in index.items(): self.ngram_store.insert([{ "ngram" : ngram, "code" : code, "count" : count } for code, count in counter.items() ]) self.ngram_store.ensure_index('ngram') def bulk_load(self, code_to_geom, geom_serializer=None): index = {} fc = { 'type' : 'GeometryCollection' } features = [] for code, geom in code_to_geom.items(): ngrams = self.parser(code) for ngram in ngrams: if not ngram in index: index[ngram] = Counter() index[ngram][code] += 1 if geom_serializer: geom = json.loads(geom_serializer(geom)) else: geom = json.loads(self.serialize(geom)) orig = code features.append({ '_id' : code, 'geometry' : geom, 'properties' : { 'name' : orig } }) log.debug('bulk_load got {n} features'.format(n=len(features))) fc['features'] = features self.code_store.insert_features(fc) for ngram, counter in index.items(): self.ngram_store.insert([{ "ngram" : ngram, "code" : code, "count" : count } for code, count in counter.items() ]) self.code_store.create_index('ngram') def bulk_geocode(self, codes, srid=None): return (self[code] for code in codes) def reverse_geocode(self, geometry): val = self.code_store.find_features(geo_query={'bboverlaps' : geometry}) try: return val.next() except StopIteration: return None def drop(self): self.code_store.drop() self.ngram_store.drop()
class ExactGeocoder(object, UserDict.DictMixin): def __init__(self, name, case_sensitive=False, long_codes=False, srid=4326, fc=None, clear=False): self.code_store = GeoJSONCollection( db=settings.MONGODB_ROUTES['ga_geocoder'] if 'ga_geocoder' in settings.MONGODB_ROUTES else settings.MONGODB_ROUTES['default'], collection=name, index_path=app_settings.GEO_INDEX_PATH, srid=srid, fc=fc, clear=clear ) self.code_store['kind'] = 'exact' self.serialize = lambda x: x.json self.deserialize = lambda x: GEOSGeometry(x, srid=self.code_store.srid) # # setup parser # case_sensitive = self.code_store['case_sensitive'] if 'case_sensitive' in self.code_store else case_sensitive long_codes = self.code_store['long_codes'] if 'long_codes' in self.code_store else long_codes self.code_store['name'] = name self.code_store['case_sensitive'] = case_sensitive self.code_store['long_codes'] = long_codes if case_sensitive and long_codes: self.parse = ci_shortcode elif case_sensitive: self.parse = ci_code elif long_codes: self.parse = cs_shortcode else: self.parse = cs_code def __setitem__(self, code, geometry): orig = code code = self.parse(code) self.code_store.insert_features({ "_id" : code, 'type' : 'Feature', 'properties' : { 'code' : code, 'name' : orig }, 'geometry' : self.serialize(geometry) }) def bulk_load(self, code_to_geom, geom_serializer=None): fc = { 'type' : "FeatureCollection", } log.debug("Beginning bulk load of {name}".format(name=self.code_store['name'])) features = [] for code, geom in code_to_geom.items(): if geom_serializer: geom = json.loads(geom_serializer(geom)) else: geom = json.loads(self.serialize(geom)) orig = code code = self.parse(code) features.append({ '_id' : code, 'geometry' : geom, 'properties' : { 'code' : code, 'name' : orig } }) log.debug('bulk_load got {n} features'.format(n=len(features))) fc['features'] = features self.code_store.insert_features(fc) def __getitem__(self, code): srid=None if isinstance(code, tuple): code, srid = code val = self.code_store.coll.find_one(code) if val: if srid: g = GEOSGeometry(json.dumps(val['geometry']), srid=self.code_store.srid) g.transform(srid) val['geometry'] = json.loads(g.json) return val else: raise KeyError(code) def __delitem__(self, code): self.code_store.coll.remove(code) def keys(self): return self.code_store.keys() def bulk_geocode(self, codes, srid=None): chunks = _chunk(codes) for chunk in chunks: features = self.code_store.find_features(spec={ "_id" : { "$in" : chunk }}) if srid: for feature in features: g = GEOSGeometry(json.dumps(feature['geometry']), srid=self.code_store.srid) g.transform(srid) feature['geometry'] = json.loads(g.json) yield feature['_id'], feature else: for feature in features: yield feature['_id'], feature def reverse_geocode(self, geometry): val = self.code_store.find_features(geo_query={'bboverlaps' : geometry}) try: return val.next() except StopIteration: return None def drop(self): self.code_store.drop() self.code_store = None
class GeoMongoTests(TestCase): def setUp(self): self.collection = GeoJSONCollection(mongo.test, 'test_collection', './test_indices', srid=4326, clear=True) self.collection.insert_features({ 'type' : "FeatureCollection", 'property3' : 3, "features" : [ { 'type' : 'Feature', 'geometry' : json.loads(poly1.json), 'properties' : {'name' : 'poly1'} }, { 'type' : 'Feature', 'geometry' : json.loads(poly2.json), 'properties' : {'name' : 'poly2'} }, { 'type' : 'Feature', 'geometry' : json.loads(poly3.json), 'properties' : {'name' : 'poly3'} } ] }) def tearDown(self): self.collection.drop() def test_load(self): """Test loading with a few features""" self.assertEquals(self.collection.count(), 3, 'count() failed basic test') self.assertEquals(self.collection.count(cover), 3, "count() failed cover test") def test_geoquery(self): result = list(self.collection.find_features(geo_spec={'$within' : cover})) self.assertEquals(len(result),3) self.assertEquals(result[0]['properties']['name'], 'poly1') self.assertEquals(result[1]['properties']['name'], 'poly2') self.assertEquals(result[2]['properties']['name'], 'poly3') def test_combined_query(self): result = list(self.collection.find_features(geo_spec={'$within' : cover}, spec={'properties.name' : 'poly1'})) self.assertEquals(len(result), 1) self.assertEquals(result[0]['properties']['name'], 'poly1') def test_plain_query(self): """Query based on properties instead of geometry""" result = list(self.collection.find_features(spec={'properties.name' : 'poly1'})) result2 = list(self.collection.find_features()) self.assertEquals(len(result), 1) self.assertEquals(result[0]['properties']['name'], 'poly1') self.assertEquals(len(result2), 3) def test_empty_resultset(self): """Make a query that comes out with nothing""" result1 = list(self.collection.find_features(spec={'properties.name' : 'nonexistent'})) result2 = list(self.collection.find_features(spec={'properties.name' : 'nonexistent'}, geo_spec={'$within' : cover})) result3 = list(self.collection.find_features(spec={'properties.name' : 'nonexistent'}, geo_spec={'$within' : poly0})) result4 = list(self.collection.find_features(geo_spec={'$within' : poly0})) self.assertEquals([], result1) self.assertEquals([], result2) self.assertEquals([], result3) self.assertEquals([], result4) def test_insert_individual_features(self): """Test the insertion of individual features instead of feature collection""" self.collection.insert_features([{ 'type' : 'Feature', 'geometry' : json.loads(poly1.json), 'properties' : {'name' : 'poly4'} }]) result = list(self.collection.find_features(spec={'properties.name' : 'poly4'})) self.assertEqual(len(result), 1) def test_properties(self): """Test loading a feature-collection with extra properties.""" self.collection['property1'] = 1 self.collection['property2'] = 2 self.assertEquals(self.collection['property1'], 1) self.assertEquals(self.collection['property2'], 2) del self.collection['property1'] self.assertFalse('property1' in self.collection.meta['properties']) self.assertFalse('property1' in self.collection) self.assertEquals(len(self.collection.keys()), 2) def test_delete_feature(self): """Test deleting a feature""" result = list(self.collection.find_features(spec={'properties.name' : 'poly1'})) [0] self.collection.delete_feature(result['_id']) def test_delete_features(self): """Test deleting some features""" self.collection.delete_features(geo_spec={"$within" : cover}) result1 = list(self.collection.find_features(geo_spec={'$within' : cover})) result2 = list(self.collection.find_features(spec={'properties.name' : 'poly1'})) self.assertEquals([], result1) self.assertEquals([], result2)