From 9a35c2e5c64e75a56221e38dee3c985ee2131543 Mon Sep 17 00:00:00 2001
From: Jeff Tratner <jeffrey.tratner@gmail.com>
Date: Tue, 5 Nov 2013 22:50:11 -0500
Subject: [PATCH] ENH: Add set_dataframe method to DataFrame

---
 geopandas/geodataframe.py  | 17 +++++++++++++++++
 tests/test_geodataframe.py | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py
index 2f1e662..3091e52 100644
--- a/geopandas/geodataframe.py
+++ b/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)
diff --git a/tests/test_geodataframe.py b/tests/test_geodataframe.py
index cacc18f..8e4486e 100644
--- a/tests/test_geodataframe.py
+++ b/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)