From 43616470e7a8d5a415b77cc4eb0c1dce9f12a0e0 Mon Sep 17 00:00:00 2001
From: Jacob Wasserman <jwasserman@gmail.com>
Date: Wed, 30 Oct 2013 02:17:27 -0400
Subject: [PATCH] ENH: Add geometry property to GeoDataFrame

Adds a geometry property and set_geometry() method to GeoDataFrame. This was
suggested in Issue #45. The implementation mimics that of DataFrame's
set_index in that it returns a copy of the GeoDataFrame by default with the
new geometry. You can specify a column or give a list/ndarray of geometries.
---
 geopandas/geodataframe.py  | 54 +++++++++++++++++++++++++++++++++++++-
 tests/test_geodataframe.py | 37 ++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py
index 48772b8..5ed55ac 100644
--- a/geopandas/geodataframe.py
+++ b/geopandas/geodataframe.py
@@ -4,7 +4,7 @@ import os
 
 import fiona
 import numpy as np
-from pandas import DataFrame
+from pandas import DataFrame, Series
 from shapely.geometry import mapping
 
 from geopandas import GeoSeries
@@ -24,6 +24,58 @@ class GeoDataFrame(DataFrame):
         super(GeoDataFrame, self).__init__(*args, **kwargs)
         self.crs = crs
 
+    @property
+    def geometry(self):
+        return self['geometry']
+
+    def set_geometry(self, col, drop=True, inplace=False):
+        """
+        Set the GeoDataFrame geometry using either an existing column or 
+        the specified input. By default yields a new object.
+
+        The original geometry column is replaced with the input.
+
+        Parameters
+        ----------
+        keys : column label or array
+        drop : boolean, default True
+            Delete column to be used as the new geometry
+        inplace : boolean, default False
+            Modify the GeoDataFrame in place (do not create a new object)
+
+        Examples
+        --------
+        >>> df1 = df.set_geometry([Point(0,0), Point(1,1), Point(2,2)])
+        >>> df2 = df.set_geometry('geom1')
+
+        Returns
+        -------
+        geodataframe : GeoDataFrame
+        """
+        # Most of the code here is taken from DataFrame.set_index()
+        if inplace:
+            frame = self
+        else:
+            frame = self.copy()
+
+        to_remove = None
+        if isinstance(col, Series):
+            level = col.values
+        elif isinstance(col, (list, np.ndarray)):
+            level = col
+        else:
+            level = frame[col].values
+            if drop:
+                to_remove = col
+
+        if to_remove:
+            del frame[to_remove]
+
+        frame['geometry'] = level
+
+        if not inplace:
+            return frame
+
     @classmethod
     def from_file(cls, filename, **kwargs):
         """
diff --git a/tests/test_geodataframe.py b/tests/test_geodataframe.py
index c9d90a6..4e5005c 100644
--- a/tests/test_geodataframe.py
+++ b/tests/test_geodataframe.py
@@ -35,6 +35,43 @@ class TestDataFrame(unittest.TestCase):
         self.assertTrue(type(self.df2) is GeoDataFrame)
         self.assertTrue(self.df2.crs == self.crs)
 
+    def test_set_geometry(self):
+        geom = [Point(x,y) for x,y in zip(range(5), range(5))]
+        df2 = self.df.set_geometry(geom)
+        self.assert_(self.df is not df2)
+        for x, y in zip(df2.geometry.values, geom):
+            self.assertEqual(x, y)
+
+    def test_set_geometry_col(self):
+        g = self.df.geometry
+        g_simplified = g.simplify(100)
+        self.df['simplified_geometry'] = g_simplified
+        df2 = self.df.set_geometry('simplified_geometry')
+
+        # Drop is true by default
+        self.assert_('simplified_geometry' not in df2)
+
+        for x, y in zip(df2.geometry.values, g_simplified):
+            self.assertEqual(x, y)
+
+    def test_set_geometry_col_no_drop(self):
+        g = self.df.geometry
+        g_simplified = g.simplify(100)
+        self.df['simplified_geometry'] = g_simplified
+        df2 = self.df.set_geometry('simplified_geometry', drop=False)
+
+        self.assert_('simplified_geometry' in df2)
+
+        for x, y in zip(df2.geometry.values, g_simplified):
+            self.assertEqual(x, y)
+
+    def test_set_geometry_inplace(self):
+        geom = [Point(x,y) for x,y in zip(range(5), range(5))]
+        ret = self.df.set_geometry(geom, inplace=True)
+        self.assert_(ret is None)
+        for x, y in zip(self.df['geometry'].values, geom):
+            self.assertEqual(x, y)
+
     def test_to_json(self):
         text = self.df.to_json()
         data = json.loads(text)