Merge pull request #39 from jwass/geocode

Initial commit of geocoding functionality

geopandas/geocode.py

@@ -0,0 +1,98 @@
+from collections import defaultdict
+
+import fiona
+import geopy
+from geopy.geocoders.base import GeocoderResultError
+import numpy as np
+import pandas as pd
+from shapely.geometry import Point
+
+import geopandas as gpd
+
+def geocode(strings, provider='googlev3', **kwargs):
+    """
+    Geocode a set of strings and get a GeoDataFrame of the resulting points.
+
+    Parameters
+    ----------
+    strings : list or Series of addresses to geocode
+    provider : geopy geocoder to use, default 'googlev3'
+        Some providers require additional arguments such as access keys
+        See each geocoder's specific parameters in geopy.geocoders
+        * googlev3, default
+        * bing
+        * google
+        * yahoo
+        * mapquest
+        * openmapquest
+    
+    Ensure proper use of the results by consulting the Terms of Service for
+    your provider.
+
+    Geocoding requires geopy. Install it using 'pip install geopy'. See also
+    https://github.com/geopy/geopy
+
+    Example
+    -------
+    >>> df = geocode(['boston, ma', '1600 pennsylvania ave. washington, dc'])
+    address                                               geometry
+    0                                    Boston, MA, USA  POINT (-71.0597731999999951 42.3584308000000007)
+    1  1600 Pennsylvania Avenue Northwest, President'...  POINT (-77.0365122999999983 38.8978377999999978)
+
+    """
+    if not isinstance(strings, pd.Series):
+        strings = pd.Series(strings)
+
+    coders = {'googlev3': geopy.geocoders.GoogleV3,
+              'bing': geopy.geocoders.Bing,
+              'google': geopy.geocoders.Google,
+              'yahoo': geopy.geocoders.Yahoo,
+              'mapquest': geopy.geocoders.MapQuest,
+              'openmapquest': geopy.geocoders.OpenMapQuest}
+
+    if provider not in coders:
+        raise ValueError('Unknown geocoding provider: {}'.format(provider))
+
+    coder = coders[provider](**kwargs)
+    results = {}
+    for i, s in strings.iteritems():
+        try:
+            results[i] = coder.geocode(s)
+        except (GeocoderResultError, ValueError):
+            results[i] = (None, None)
+
+    df = _prepare_geocode_result(results)
+    return df
+
+def _prepare_geocode_result(results):
+    """
+    Helper function for the geocode function
+
+    Takes a dict where keys are index entries, values are tuples containing:
+    (address, (lat, lon))
+
+    """
+    # Prepare the data for the DataFrame as a dict of lists
+    d = defaultdict(list)
+    index = []
+
+    for i, s in results.iteritems():
+        address, loc = s
+
+        # loc is lat, lon and we want lon, lat
+        if loc is None:
+            p = Point()
+        else:
+            p = Point(loc[1], loc[0])
+
+        if address is None:
+            address = pd.np.nan
+
+        d['geometry'].append(p)
+        d['address'].append(address)
+        index.append(i)
+
+    df = gpd.GeoDataFrame(d, index=index)
+    df.crs = fiona.crs.from_epsg(4326)
+
+    return df

requirements.test.txt

@@ -1 +1,2 @@
 psycopg2>=2.5.1
+geopy>=0.95.1

tests/test_geocode.py

@@ -0,0 +1,52 @@
+import unittest
+
+import fiona
+import pandas as pd
+from shapely.geometry import Point
+import geopandas as gpd
+
+from geopandas.geocode import geocode, _prepare_geocode_result
+
+class TestGeocode(unittest.TestCase):
+    def test_prepare_result(self):
+        # Calls _prepare_result with sample results from the geocoder call
+        # loop
+        p0 = Point(12.3, -45.6) # Treat these as lat/lon
+        p1 = Point(-23.4, 56.7)
+        d = {'a': ('address0', p0.coords[0]),
+             'b': ('address1', p1.coords[0])}
+
+        df = _prepare_geocode_result(d)
+        assert type(df) is gpd.GeoDataFrame
+        self.assertEqual(fiona.crs.from_epsg(4326), df.crs)
+        self.assertEqual(len(df), 2)
+        self.assert_('address' in df)
+
+        coords = df.loc['a']['geometry'].coords[0]
+        test = p0.coords[0]
+        # Output from the df should be lon/lat
+        self.assertAlmostEqual(coords[0], test[1])
+        self.assertAlmostEqual(coords[1], test[0])
+
+        coords = df.loc['b']['geometry'].coords[0]
+        test = p1.coords[0]
+        self.assertAlmostEqual(coords[0], test[1])
+        self.assertAlmostEqual(coords[1], test[0])
+
+    def test_prepare_result_none(self):
+        p0 = Point(12.3, -45.6) # Treat these as lat/lon
+        d = {'a': ('address0', p0.coords[0]),
+             'b': (None, None)}
+
+        df = _prepare_geocode_result(d)
+        assert type(df) is gpd.GeoDataFrame
+        self.assertEqual(fiona.crs.from_epsg(4326), df.crs)
+        self.assertEqual(len(df), 2)
+        self.assert_('address' in df)
+
+        row = df.loc['b']
+        self.assertEqual(len(row['geometry'].coords), 0)
+        self.assert_(pd.np.isnan(row['address']))
+    
+    def test_bad_provider(self):
+        self.assertRaises(ValueError, geocode, ['cambridge, ma'], 'badprovider')