diff --git a/pygmt/src/choropleth.py b/pygmt/src/choropleth.py index 09f77f2bae1..0e18e3f6d36 100644 --- a/pygmt/src/choropleth.py +++ b/pygmt/src/choropleth.py @@ -8,7 +8,8 @@ from pygmt._typing import GeoLike, PathLike from pygmt.alias import Alias, AliasSystem from pygmt.clib import Session -from pygmt.helpers import build_arg_list, fmt_docstring +from pygmt.exceptions import GMTTypeError, GMTValueError +from pygmt.helpers import build_arg_list, data_kind, fmt_docstring from pygmt.params import Axis, Frame __doctest_skip__ = ["choropleth"] @@ -101,6 +102,21 @@ def choropleth( >>> fig.colorbar(frame=True) >>> fig.show() """ + # The input data must be a geo-like object or a path to an OGR_GMT file. + if (kind := data_kind(data)) not in {"geojson", "file"}: + raise GMTTypeError( + type(data), + reason="Expected a geo-like object or a path to an OGR_GMT file.", + ) + + # Check if the column name exists in the data for geo-like objects. + if kind == "geojson" and (_columns := getattr(data, "columns", None)) is not None: + # The geometry column is not an aspatial attribute field, so exclude it. + geometry = getattr(getattr(data, "geometry", None), "name", None) + fields = [str(col) for col in _columns if col != geometry] + if column not in fields: + raise GMTValueError(column, description="column name", choices=fields) + aliasdict = AliasSystem( C=Alias(cmap, name="cmap"), I=Alias(intensity, name="intensity"), diff --git a/pygmt/tests/test_choropleth.py b/pygmt/tests/test_choropleth.py index 2cc9502bac5..379a2782195 100644 --- a/pygmt/tests/test_choropleth.py +++ b/pygmt/tests/test_choropleth.py @@ -2,8 +2,10 @@ Test Figure.choropleth. """ +import numpy as np import pytest from pygmt import Figure, makecpt +from pygmt.exceptions import GMTTypeError, GMTValueError geopandas = pytest.importorskip("geopandas") @@ -31,3 +33,30 @@ def test_choropleth(world): fig.choropleth(world, column="POP_EST", pen="0.3p,gray10") fig.colorbar(frame=True) return fig + + +def test_choropleth_invalid_column(world): + """ + Test that a nonexistent column raises an error. + """ + fig = Figure() + with pytest.raises(GMTValueError): + fig.choropleth(world, column="invalid") + + +def test_choropleth_geometry_column(world): + """ + Test that the geometry column is rejected, since it's not an attribute field. + """ + fig = Figure() + with pytest.raises(GMTValueError, match="Invalid column name: 'geometry'"): + fig.choropleth(world, column="geometry") + + +def test_choropleth_invalid_data_kind(): + """ + Test that data that is neither geo-like nor a file name raises an error. + """ + fig = Figure() + with pytest.raises(GMTTypeError, match="Unrecognized data type"): + fig.choropleth(np.array([[1.0, 2.0], [3.0, 4.0]]), column="POP")