From cb0c8b5ee4cf6da8e1fc70ffbd23e486f1919724 Mon Sep 17 00:00:00 2001 From: Matt Richards <45483497+m-richards@users.noreply.github.com> Date: Tue, 22 Feb 2022 17:33:36 +1000 Subject: [PATCH] API: improve handling of invalid geo column for .geometry (more informative error message) + update getitem (preserve GeoDataFrame if a geometry dtype is present) (#2329) Co-authored-by: Joris Van den Bossche Co-authored-by: Martin Fleischmann --- geopandas/geodataframe.py | 50 +++++++++++++++++++------ geopandas/tests/test_geodataframe.py | 22 +++++++++++ geopandas/tests/test_op_output_types.py | 39 ++++++++++++++----- 3 files changed, 91 insertions(+), 20 deletions(-) diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py index 4127003..3b4b516 100644 --- a/geopandas/geodataframe.py +++ b/geopandas/geodataframe.py @@ -205,10 +205,33 @@ class GeoDataFrame(GeoPandasBase, DataFrame): 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 - ) + if self._geometry_column_name is None: + msg = ( + "You are calling a geospatial method on the GeoDataFrame, " + "but the active geometry column to use has not been set. " + ) + else: + msg = ( + "You are calling a geospatial method on the GeoDataFrame, " + f"but the active geometry column ('{self._geometry_column_name}') " + "is not present. " + ) + geo_cols = list(self.columns[self.dtypes == "geometry"]) + if len(geo_cols) > 0: + msg += ( + f"\nThere are columns with geometry data type ({geo_cols}), and " + "you can either set one as the active geometry with " + 'df.set_geometry("name") or access the column as a ' + 'GeoSeries (df["name"]) and call the method directly on it.' + ) + else: + msg += ( + "\nThere are no existing columns with geometry data type. You can " + "add a geometry column as the active geometry column with " + "df.set_geometry. " + ) + + raise AttributeError(msg) return self[self._geometry_column_name] def _set_geometry(self, col): @@ -1329,18 +1352,23 @@ individually so that features may have different properties def __getitem__(self, key): """ If the result is a column containing only 'geometry', return a - GeoSeries. If it's a DataFrame with a 'geometry' column, return a - GeoDataFrame. + GeoSeries. If it's a DataFrame with any columns of GeometryDtype, + return a GeoDataFrame. """ result = super().__getitem__(key) geo_col = self._geometry_column_name if isinstance(result, Series) and isinstance(result.dtype, GeometryDtype): result.__class__ = GeoSeries - elif isinstance(result, DataFrame) and geo_col in result: - result.__class__ = GeoDataFrame - result._geometry_column_name = geo_col - elif isinstance(result, DataFrame) and geo_col not in result: - result.__class__ = DataFrame + elif isinstance(result, DataFrame): + if (result.dtypes == "geometry").sum() > 0: + result.__class__ = GeoDataFrame + if geo_col in result: + result._geometry_column_name = geo_col + else: + result._geometry_column_name = None + result._crs = None + else: + result.__class__ = DataFrame return result def __setitem__(self, key, value): diff --git a/geopandas/tests/test_geodataframe.py b/geopandas/tests/test_geodataframe.py index c81a821..630021d 100644 --- a/geopandas/tests/test_geodataframe.py +++ b/geopandas/tests/test_geodataframe.py @@ -334,6 +334,28 @@ class TestDataFrame: assert isinstance(result, GeoDataFrame) assert isinstance(result.index, pd.DatetimeIndex) + def test_get_geometry_invalid(self): + df = GeoDataFrame() + df["geom"] = self.df.geometry + msg_geo_col_none = "active geometry column to use has not been set. " + msg_geo_col_missing = "is not present. " + + with pytest.raises(AttributeError, match=msg_geo_col_missing): + df.geometry + df2 = self.df.copy() + df2["geom2"] = df2.geometry + df2 = df2[["BoroCode", "BoroName", "geom2"]] + with pytest.raises(AttributeError, match=msg_geo_col_none): + df2.geometry + + msg_other_geo_cols_present = "There are columns with geometry data type" + msg_no_other_geo_cols = "There are no existing columns with geometry data type" + with pytest.raises(AttributeError, match=msg_other_geo_cols_present): + df2.geometry + + with pytest.raises(AttributeError, match=msg_no_other_geo_cols): + GeoDataFrame().geometry + def test_align(self): df = self.df2 diff --git a/geopandas/tests/test_op_output_types.py b/geopandas/tests/test_op_output_types.py index 3be6b35..291abfb 100644 --- a/geopandas/tests/test_op_output_types.py +++ b/geopandas/tests/test_op_output_types.py @@ -60,7 +60,18 @@ def assert_object(result, expected_type, geo_name="geometry", crs=crs_wgs): if geo_name is not None: _check_metadata_gdf(result, geo_name=geo_name, crs=crs) else: - with pytest.raises(AttributeError, match="No geometry data set yet"): + if result._geometry_column_name is None: + msg = ( + "You are calling a geospatial method on the GeoDataFrame, " + "but the active" + ) + else: + msg = ( + "You are calling a geospatial method on the GeoDataFrame, but " + r"the active geometry column \(" + rf"'{result._geometry_column_name}'\) is not present" + ) + with pytest.raises(AttributeError, match=msg): result.geometry.name # be explicit that geometry is invalid here elif expected_type == GeoSeries: _check_metadata_gs(result, name=geo_name, crs=crs) @@ -71,8 +82,8 @@ def test_getitem(df): assert_object(df[["value1", "value2"]], pd.DataFrame) assert_object(df[[geo_name, "geometry2"]], GeoDataFrame, geo_name) assert_object(df[[geo_name]], GeoDataFrame, geo_name) - assert_object(df[["geometry2", "value1"]], pd.DataFrame) - assert_object(df[["geometry2"]], pd.DataFrame) + assert_object(df[["geometry2", "value1"]], GeoDataFrame, None, None) + assert_object(df[["geometry2"]], GeoDataFrame, None, None) assert_object(df[["value1"]], pd.DataFrame) # Series assert_object(df[geo_name], GeoSeries, geo_name) @@ -119,9 +130,7 @@ def test_iloc(df): def test_squeeze(df): geo_name = df.geometry.name assert_object(df[[geo_name]].squeeze(), GeoSeries, geo_name) - - # Not ideal behaviour, but this is consistent with __getitem__ - assert_object(df[["geometry2"]].squeeze(), pd.Series) + assert_object(df[["geometry2"]].squeeze(), GeoSeries, "geometry2", crs=crs_osgb) def test_to_frame(df): @@ -192,8 +201,8 @@ def test_apply(df): assert_object(df[["value1", "value2"]].apply(identity), pd.DataFrame) assert_object(df[[geo_name, "geometry2"]].apply(identity), GeoDataFrame, geo_name) assert_object(df[[geo_name]].apply(identity), GeoDataFrame, geo_name) - assert_object(df[["geometry2", "value1"]].apply(identity), pd.DataFrame) - assert_object(df[["geometry2"]].apply(identity), pd.DataFrame) + assert_object(df[["geometry2", "value1"]].apply(identity), GeoDataFrame, None, None) + assert_object(df[["geometry2"]].apply(identity), GeoDataFrame, None, None) assert_object(df[["value1"]].apply(identity), pd.DataFrame) # axis = 0, Series @@ -211,6 +220,18 @@ def test_apply(df): df[[geo_name, "geometry2"]].apply(identity, axis=1), GeoDataFrame, geo_name ) assert_object(df[[geo_name]].apply(identity, axis=1), GeoDataFrame, geo_name) + # TODO below should be a GeoDataFrame to be consistent with new getitem logic + # leave as follow up as quite complicated + # FrameColumnApply.series_generator returns object dtypes Series, so will have + # patch result of apply assert_object(df[["geometry2", "value1"]].apply(identity, axis=1), pd.DataFrame) - assert_object(df[["geometry2"]].apply(identity, axis=1), pd.DataFrame) + assert_object(df[["value1"]].apply(identity, axis=1), pd.DataFrame) + # if compat # https://github.com/pandas-dev/pandas/pull/30091 + + +def test_apply_axis1_secondary_geo_cols(df): + def identity(x): + return x + + assert_object(df[["geometry2"]].apply(identity, axis=1), GeoDataFrame, None, None)