Integration with geopandas #588

.gitignore

@@ -62,6 +62,7 @@ target/
 
 .ipynb_checkpoints
 .idea/*
+.vscode/*
 tools/_build
 Untitled*.ipynb
 .mypy*

altair/utils/core.py

@@ -227,9 +227,13 @@ def parse_shorthand(shorthand, data=None):
     {'aggregate': 'count', 'type': 'quantitative'}
     """
     attrs = _parse_shorthand(shorthand)
-    if isinstance(data, pd.DataFrame) and 'type' not in attrs:
+    if isinstance(data, pd.DataFrame):
         if 'field' in attrs and attrs['field'] in data.columns:
-            attrs['type'] = infer_vegalite_type(data[attrs['field']])
+                if 'type' not in attrs:
+                    attrs['type'] = infer_vegalite_type(data[attrs['field']])
+                if hasattr(data,'__geo_interface__'): #TODO: Add descripion in shorthand
+                    attrs['field']='properties.'+attrs['field']
+
     return attrs
 
 

altair/utils/data.py

@@ -85,9 +85,17 @@ def to_json(data, prefix='altair-data'):
     check_data_type(data)
     ext = '.json'
     filename = _compute_filename(prefix=prefix, ext=ext)
+    data_format = {'type': 'json'}
     if isinstance(data, pd.DataFrame):
         data = sanitize_dataframe(data)
-        data.to_json(filename, orient='records')
+        if not hasattr(data,'__geo_interface__'):
+            data.to_json(filename, orient='records')
+        else: #GeoPandas
+            with open(filename) as f:
+                json.dump(data.__geo_interface__, f)
+            data_format['property']='features'
+
+
     elif isinstance(data, dict):
         if 'values' not in data:
             raise KeyError('values expected in data dict, but not present.')
@@ -96,9 +104,18 @@ def to_json(data, prefix='altair-data'):
             json.dump(values, f)
     return {
         'url': filename,
-        'format': {'type': 'json'}
+        'format': data_format
     }
-
+@curry
+def to_geojson_values(data, feature="features"):
+    if not hasattr(data, '__geo_interface__'):
+        raise TypeError('Expected GeoDataFrame or __geo_interface__, got: {}'.format(type(data)))
+    if isinstance(data, pd.DataFrame):
+        data = sanitize_dataframe(data)
+    return {
+            'values':data.__geo_interface__,
+             'format':{'type':'json','property':feature}
+            }
 
 @curry
 def to_csv(data, prefix='altair-data'):
@@ -123,6 +140,12 @@ def to_values(data):
     check_data_type(data)
     if isinstance(data, pd.DataFrame):
         data = sanitize_dataframe(data)
+        if hasattr(data,'__geo_interface__'):#GeoPandas
+            return {
+                    'values':data.__geo_interface__,
+                    'format':{'type':'json','property':'features'}
+                    }
+
         return {'values': data.to_dict(orient='records')}
     elif isinstance(data, dict):
         if 'values' not in data:

altair/utils/tests/test_data.py

@@ -2,7 +2,7 @@
 import pandas as pd
 
 
-from ..data import limit_rows, MaxRowsError, sample, pipe, to_values
+from ..data import limit_rows, MaxRowsError, sample, pipe, to_values,to_geojson_values
 
 
 def _create_dataframe(N):
@@ -63,3 +63,13 @@ def test_type_error():
     for f in (sample, limit_rows, to_values):
         with pytest.raises(TypeError):
             pipe(0, f)
+
+
+def test_to_geojson_values():
+    gpd = pytest.importorskip('geopandas')
+    geo_data = gpd.GeoDataFrame({ "name": ['a','b']},
+                        geometry=[gpd.geoseries.Point((1.0, 0.0)),
+                                  gpd.geoseries.Point((0.0, 1.0))],index=['i','j'])
+    result = pipe(geo_data, to_geojson_values)
+    assert result=={'format': {'property': 'features', 'type': 'json'},
+                        'values': geo_data.__geo_interface__}

altair/vegalite/data.py

@@ -2,7 +2,7 @@
 from toolz.curried import curry, pipe
 from ..utils.core import sanitize_dataframe
 from ..utils.data import (
-    MaxRowsError, limit_rows, sample, to_csv, to_json, to_values,
+    MaxRowsError, limit_rows, sample, to_csv, to_json, to_values, to_geojson_values,
     check_data_type, DataTransformerRegistry
 )
 
@@ -24,5 +24,6 @@ def default_data_transformer(data, max_rows=5000):
     'to_csv',
     'to_json',
     'to_values',
-    'check_data_type'
+    'check_data_type',
+    'to_geojson_values'
 )

altair/vegalite/v2/__init__.py

@@ -17,7 +17,7 @@
 from .data import (
     MaxRowsError,
     pipe, curry, limit_rows,
-    sample, to_json, to_csv, to_values,
+    sample, to_json, to_csv, to_values, to_geojson_values,
     default_data_transformer,
     data_transformers
 )

altair/vegalite/v2/data.py

@@ -1,5 +1,5 @@
 from ..data import (MaxRowsError, curry, default_data_transformer, limit_rows,
-                    pipe, sample, to_csv, to_json, to_values, DataTransformerRegistry)
+                    pipe, sample, to_csv, to_json, to_values,to_geojson_values, DataTransformerRegistry)
 
 
 # ==============================================================================
@@ -27,5 +27,6 @@
     'to_csv',
     'to_json',
     'to_values',
-    'data_transformers'
+    'data_transformers',
+    'to_geojson_values'
 )