Merge pull request #15 from cfarmer/writing

adds `to_file` method for geodataframe

geopandas/geodataframe.py

@@ -1,5 +1,5 @@
-from collections import defaultdict
-import json
+from collections import defaultdict, OrderedDict
+import json, os
 
 import fiona
 from pandas import DataFrame
@@ -8,6 +8,8 @@ from shapely.geometry import mapping, shape
 from geopandas import GeoSeries
 from plotting import plot_dataframe
 
+import numpy as np
+
 
 class GeoDataFrame(DataFrame):
     """
@@ -70,6 +72,55 @@ class GeoDataFrame(DataFrame):
             {'type': 'FeatureCollection',
              'features': [feature(i, row) for i, row in self.iterrows()]},
             **kwargs )
+            
+    def to_file(self, filename, driver="ESRI Shapefile", **kwargs):
+        """
+        Write this GeoDataFrame to an OGR data source
+        
+        A dictionary of supported OGR providers is available via:
+        >>> import fiona
+        >>> fiona.supported_drivers
+
+        Parameters
+        ----------
+        filename : string 
+            File path or file handle to write to.
+        driver : string, default 'ESRI Shapefile'
+            The OGR format driver used to write the vector file.
+
+        The *kwargs* are passed to fiona.open and can be used to write 
+        to multi-layer data, store data within archives (zip files), etc.
+        """
+        def convert_type(in_type):
+            if in_type == object:
+                return 'str'
+            return type(np.asscalar(np.zeros(1, in_type))).__name__
+            
+        def feature(i, row):
+            return {
+                'id': str(i),
+                'type': 'Feature',
+                'properties': {
+                    k: v for k, v in row.iteritems() if k != 'geometry'},
+                'geometry': mapping(row['geometry']) }
+        
+        properties = OrderedDict([(col, convert_type(_type)) for col, _type 
+            in zip(self.columns, self.dtypes) if col!='geometry'])
+        # Need to check geom_types before we write to file... 
+        # Some (most?) providers expect a single geometry type: 
+        # Point, LineString, or Polygon
+        geom_types = self['geometry'].geom_type.unique()
+        from os.path import commonprefix # To find longest common prefix
+        geom_type = commonprefix([g[::-1] for g in geom_types])[::-1]  # Reverse
+        if geom_type == '': # No common suffix = mixed geometry types
+            raise ValueError("Geometry column cannot contains mutiple "
+                             "geometry types when writing to file.")
+        schema = {'geometry': geom_type, 'properties': properties}
+        filename = os.path.abspath(os.path.expanduser(filename))
+        with fiona.open(filename, 'w', driver=driver, crs=self.crs, 
+                        schema=schema, **kwargs) as c:
+            for i, row in self.iterrows():
+                c.write(feature(i, row))
 
     def to_crs(self, crs=None, epsg=None, inplace=False):
         """Transform geometries to a new coordinate reference system

geopandas/geoseries.py

@@ -46,10 +46,26 @@ class GeoSeries(Series):
         self.crs = crs
 
     @classmethod
-    def from_file(cls, filename):
-        """Alternate constructor to create a GeoSeries from a file"""
+    def from_file(cls, filename, **kwargs):
+        """
+        Alternate constructor to create a GeoSeries from a file
+        
+        Parameters
+        ----------
+        
+        filename : str
+            File path or file handle to read from. Depending on which kwargs
+            are included, the content of filename may vary, see:
+            http://toblerity.github.io/fiona/README.html#usage
+            for usage details.
+        kwargs : key-word arguments
+            These arguments are passed to fiona.open, and can be used to 
+            access multi-layer data, data stored within archives (zip files),
+            etc.
+        
+        """
         geoms = []
-        with fiona.open(filename) as f:
+        with fiona.open(filename, **kwargs) as f:
             crs = f.crs
             for rec in f:
                 geoms.append(shape(rec['geometry']))
@@ -57,6 +73,14 @@ class GeoSeries(Series):
         g.crs = crs
         return g
 
+    def to_file(self, filename, driver="ESRI Shapefile", **kwargs):
+        from geopandas import GeoDataFrame
+        data = GeoDataFrame({"geometry": self,
+                          "id":self.index.values},
+                          index=self.index)
+        data.crs = self.crs
+        data.to_file(filename, driver, **kwargs)
+        
     #
     # Internal methods
     #

tests/test_geodataframe.py

@@ -2,6 +2,8 @@ import unittest
 import json
 import numpy as np
 
+import tempfile
+
 from geopandas import GeoDataFrame
 
 class TestDataFrame(unittest.TestCase):
@@ -23,3 +25,13 @@ class TestDataFrame(unittest.TestCase):
         data = json.loads(text)
         self.assertTrue(data['type'] == 'FeatureCollection')
         self.assertTrue(len(data['features']) == 5)
+        
+    def test_to_file(self):
+        with tempfile.NamedTemporaryFile(suffix='.shp') as t:
+            self.df.to_file(t.name)
+            # Read layer back in?
+            df = GeoDataFrame.from_file(t.name)
+            self.assertTrue('geometry' in df)
+            self.assertTrue(len(df) == 5)
+            self.assertTrue(np.alltrue(df['BoroName'].values == np.array(['Staten Island',
+                                   'Queens', 'Brooklyn', 'Manhattan', 'Bronx'])))