ENH: Add set_dataframe method to DataFrame

geopandas/geodataframe.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 import json
 import os
+import sys
 
 import fiona
 import numpy as np
@@ -13,6 +14,7 @@ import geopandas.io
 
 
 DEFAULT_GEO_COLUMN_NAME = 'geometry'
+PY3 = sys.version[0] == 3
 
 
 class GeoDataFrame(DataFrame):
@@ -332,3 +334,18 @@ class GeoDataFrame(DataFrame):
 
     def plot(self, *args, **kwargs):
         return plot_dataframe(self, *args, **kwargs)
+
+def _dataframe_set_geometry(self, col, drop=False, inplace=False, crs=None):
+    if inplace:
+        raise ValueError("Can't do inplace setting when converting from"
+                         " DataFrame to GeoDataFrame")
+    gf = GeoDataFrame(self)
+    # this will copy so that BlockManager gets copied
+    return gf.set_geometry(col, drop=drop, inplace=False, crs=crs)
+
+if PY3:
+    DataFrame.set_geometry = _dataframe_set_geometry
+else:
+    import types
+    DataFrame.set_geometry = types.MethodType(_dataframe_set_geometry, None,
+                                              DataFrame)

tests/test_geodataframe.py

@@ -5,6 +5,7 @@ import tempfile
 import shutil
 
 import numpy as np
+import pandas as pd
 from shapely.geometry import Point, Polygon
 
 
@@ -55,7 +56,9 @@ class TestDataFrame(unittest.TestCase):
         # reset so it outputs okay
         df2.crs = df.crs
         tu.assert_geoseries_equal(df2.geometry, GeoSeries(geom2))
-        tu.assert_geoseries_equal(df2['location'], df['location'])
+        # for right now, non-geometry comes back as series
+        tu.assert_geoseries_equal(df2['location'], df['location'],
+                                  check_series_type=False, check_dtype=False)
 
     def test_geometry_property(self):
         tu.assert_geoseries_equal(self.df.geometry, self.df['geometry'],
@@ -270,3 +273,29 @@ class TestDataFrame(unittest.TestCase):
             con.close()
 
         tu.validate_boro_df(self, df)
+
+    def test_dataframe_to_geodataframe(self):
+        df = pd.DataFrame({"A": range(len(self.df)), "location":
+                           list(self.df.geometry)}, index=self.df.index)
+        gf = df.set_geometry('location', crs=self.df.crs)
+        tu.assert_isinstance(df, pd.DataFrame)
+        tu.assert_isinstance(gf, GeoDataFrame)
+        tu.assert_geoseries_equal(gf.geometry, self.df.geometry)
+        self.assertEqual(gf.geometry.name, 'location')
+        self.assert_('geometry' not in gf)
+
+        gf2 = df.set_geometry('location', crs=self.df.crs, drop=True)
+        tu.assert_isinstance(df, pd.DataFrame)
+        tu.assert_isinstance(gf2, GeoDataFrame)
+        self.assertEqual(gf2.geometry.name, 'geometry')
+        self.assert_('geometry' in gf2)
+        self.assert_('location' not in gf2)
+        self.assert_('location' in df)
+
+        # should be a copy
+        df.ix[0, "A"] = 100
+        self.assertEqual(gf.ix[0, "A"], 0)
+        self.assertEqual(gf2.ix[0, "A"], 0)
+
+        self.assertRaises(ValueError, df.set_geometry, 'location',
+                          inplace=True)