From 7df6db234fd7ac432f7ea8c1fbb76685f5baffeb Mon Sep 17 00:00:00 2001 From: Jeff Tratner Date: Sun, 3 Nov 2013 13:34:08 -0500 Subject: [PATCH] ENH: Allow geometry property to be set. * Allows for clearer errors when a GeoDataFrame is created but has no geometry. * TST: Add test util for comparing sequence (b/c assert_almost_equal seems to not work with shapely. --- examples/nyc_boros.py | 2 +- geopandas/geodataframe.py | 24 +++++++++++++++++---- geopandas/plotting.py | 4 ++-- tests/test_geodataframe.py | 43 ++++++++++++++++++++++++++++++-------- tests/test_types.py | 2 +- tests/util.py | 11 +++++++++- 6 files changed, 68 insertions(+), 18 deletions(-) diff --git a/examples/nyc_boros.py b/examples/nyc_boros.py index 744664d..e2ebc14 100644 --- a/examples/nyc_boros.py +++ b/examples/nyc_boros.py @@ -22,7 +22,7 @@ boros.plot() plt.xticks(rotation=90) plt.savefig('nyc.png', dpi=DPI, bbox_inches='tight') #plt.show() -boros['geometry'].convex_hull.plot() +boros.geometry.convex_hull.plot() plt.xticks(rotation=90) plt.savefig('nyc_hull.png', dpi=DPI, bbox_inches='tight') #plt.show() diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py index ed9be90..5c31afb 100644 --- a/geopandas/geodataframe.py +++ b/geopandas/geodataframe.py @@ -17,16 +17,32 @@ class GeoDataFrame(DataFrame): A GeoDataFrame object is a pandas.DataFrame that has a column named 'geometry' which is a GeoSeries. """ - _metadata = ['crs'] + _metadata = ['crs', '_geometry_column_name'] + _geometry_column_name = 'geometry' def __init__(self, *args, **kwargs): crs = kwargs.pop('crs', None) super(GeoDataFrame, self).__init__(*args, **kwargs) self.crs = crs - @property - def geometry(self): - return self['geometry'] + def _get_geometry(self): + if self._geometry_column_name not in self: + raise AttributeError("No geometry data set yet (expected in" + " column '%s'." % self._geometry_column_name) + return self[self._geometry_column_name] + + def _set_geometry(self, col): + try: + if col in self: + raise ValueError("Can't set a column name via the geometry" + " property") + except TypeError: # hashing issues + pass + + self.set_geometry(self, col, inplace=True) + + geometry = property(fget=_get_geometry, fset=_set_geometry, + doc="Geometry data for GeoDataFrame") def set_geometry(self, col, drop=True, inplace=False): """ diff --git a/geopandas/plotting.py b/geopandas/plotting.py index 22e67ca..f670eed 100644 --- a/geopandas/plotting.py +++ b/geopandas/plotting.py @@ -163,7 +163,7 @@ def plot_dataframe(s, column=None, colormap=None, alpha=0.5, from matplotlib.colors import Normalize from matplotlib import cm if column is None: - return plot_series(s['geometry'], colormap=colormap, alpha=alpha, axes=axes) + return plot_series(s.geometry, colormap=colormap, alpha=alpha, axes=axes) else: if s[column].dtype is np.dtype('O'): categorical = True @@ -185,7 +185,7 @@ def plot_dataframe(s, column=None, colormap=None, alpha=0.5, ax = plt.gca() else: ax = axes - for geom, value in zip(s['geometry'], values): + for geom, value in zip(s.geometry, values): if geom.type == 'Polygon' or geom.type == 'MultiPolygon': plot_multipolygon(ax, geom, facecolor=cmap.to_rgba(value), alpha=alpha) elif geom.type == 'LineString' or geom.type == 'MultiLineString': diff --git a/tests/test_geodataframe.py b/tests/test_geodataframe.py index ca69057..cd423e5 100644 --- a/tests/test_geodataframe.py +++ b/tests/test_geodataframe.py @@ -35,12 +35,40 @@ class TestDataFrame(unittest.TestCase): self.assertTrue(type(self.df2) is GeoDataFrame) self.assertTrue(self.df2.crs == self.crs) + def test_geometry_property(self): + tests.util.assert_seq_equal(self.df.geometry, self.df['geometry']) + df = self.df.copy() + new_geom = [Point(x,y) for x, y in zip(range(len(self.df)), + range(len(self.df)))] + + df.geometry = new_geom + tests.util.assert_seq_equal(df.geometry, new_geom) + # should this be tested here? + tests.util.assert_seq_equal(df['geometry'], new_geom) + + def _should_raise_att_error(): + df = self.df.copy() + del df['geometry'] + df.geometry + + self.assertRaises(AttributeError, _should_raise_att_error) + + def _should_raise_key_error(): + df = self.df.copy() + del df['geometry'] + df['geometry'] + + self.assertRaises(KeyError, _should_raise_key_error) + def test_set_geometry(self): geom = [Point(x,y) for x,y in zip(range(5), range(5))] + original_geom = self.df.geometry + 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) + tests.util.assert_seq_equal(df2.geometry, geom) + tests.util.assert_seq_equal(self.df.geometry, original_geom) + tests.util.assert_seq_equal(self.df['geometry'], self.df.geometry) def test_set_geometry_col(self): g = self.df.geometry @@ -51,8 +79,7 @@ class TestDataFrame(unittest.TestCase): # 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) + tests.util.assert_seq_equal(df2.geometry, g_simplified) def test_set_geometry_col_no_drop(self): g = self.df.geometry @@ -61,16 +88,14 @@ class TestDataFrame(unittest.TestCase): 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) + tests.util.assert_seq_equal(df2.geometry, g_simplified) 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) + tests.util.assert_seq_equal(self.df['geometry'], geom) + tests.util.assert_seq_equal(self.df.geometry, geom) def test_to_json(self): text = self.df.to_json() diff --git a/tests/test_types.py b/tests/test_types.py index 551efaf..6e58c5b 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -65,7 +65,7 @@ class TestDataFrame(unittest.TestCase): for x, y in zip(range(N), range(N))]) def test_geometry(self): - assert type(self.df['geometry']) is GeoSeries + assert type(self.df.geometry) is GeoSeries def test_nongeometry(self): assert type(self.df['value1']) is Series diff --git a/tests/util.py b/tests/util.py index 16538f2..7110c0e 100644 --- a/tests/util.py +++ b/tests/util.py @@ -30,7 +30,7 @@ def validate_boro_df(test, df): columns = ('borocode', 'boroname', 'shape_leng', 'shape_area') for col in columns: test.assertTrue(col in df.columns, 'Column {} missing'.format(col)) - test.assertTrue(all(df['geometry'].type == 'MultiPolygon')) + test.assertTrue(all(df.geometry.type == 'MultiPolygon')) def connect(dbname): try: @@ -85,3 +85,12 @@ def create_db(df): con.close() return True + + +def assert_seq_equal(left, right): + """Poor man's version of assert_almost_equal which isn't working with Shapely + objects right now""" + assert len(left) == len(right), "Mismatched lengths: %d != %d" % (len(left), len(right)) + for elem_left, elem_right in zip(left, right): + assert elem_left == elem_right, "%r != %r" % (left, right) + return True