diff --git a/datawrapper/__main__.py b/datawrapper/__main__.py index 7c7538f..27125d9 100644 --- a/datawrapper/__main__.py +++ b/datawrapper/__main__.py @@ -44,6 +44,7 @@ class Datawrapper: _BASE_URL = "https://api.datawrapper.de" _CHARTS_URL = _BASE_URL + "/v3/charts" _PUBLISH_URL = _BASE_URL + "/charts" + _BASEMAPS_URL = _BASE_URL + "/v3/basemaps" _FOLDERS_URL = _BASE_URL + "/v3/folders" _ACCESS_TOKEN = os.getenv("DATAWRAPPER_ACCESS_TOKEN") @@ -625,6 +626,90 @@ def export_chart( logger.error(msg) raise Exception(msg) + def get_basemaps(self) -> list[dict[str, Any]]: + """Get a list of the available basemaps. + + Returns + ------- + list[dict] + A list of dictionaries containing the basemaps available in your Datawrapper account. + """ + _header = self._auth_header + _header["accept"] = "*/*" + + response = r.get( + url=self._BASEMAPS_URL, + headers=_header, + ) + + if response.ok: + return response.json() + else: + msg = "Couldn't retrieve basemaps in your account." + logger.error(msg) + raise Exception(msg) + + def get_basemap(self, basemap_id: str, wgs84: bool = False) -> dict[str, Any]: + """Get the metdata of the requested basemap. + + Parameters + ---------- + basemap_id : str + ID of basemap to get. + wgs84 : bool, optional + Whether to return the basemap in the WGS84 project, by default False + + Returns + ------- + dict + A dictionary containing the requested basemap's metadata. + """ + _header = self._auth_header + _header["accept"] = "*/*" + + response = r.get( + url=f"{self._BASEMAPS_URL}/{basemap_id}", + headers=_header, + params={"wgs84": wgs84}, + ) + + if response.ok: + return response.json() + else: + msg = "Couldn't retrieve basemap in your account." + logger.error(msg) + raise Exception(msg) + + def get_basemap_key(self, basemap_id: str, basemap_key: str) -> dict[str, Any]: + """Get the list of available values for a basemap's key. + + Parameters + ---------- + basemap_id : str + ID of basemap to get. + basemap_key : str + Metadata key of basemap to get. + + Returns + ------- + dict + A dictionary containing the requested data. + """ + _header = self._auth_header + _header["accept"] = "*/*" + + response = r.get( + url=f"{self._BASEMAPS_URL}/{basemap_id}/{basemap_key}", + headers=_header, + ) + + if response.ok: + return response.json() + else: + msg = "Couldn't retrieve basemap key in your account." + logger.error(msg) + raise Exception(msg) + def get_folders(self) -> dict[Any, Any] | None | Any: """Get a list of folders in your Datawrapper account. diff --git a/tests/test_basemaps.py b/tests/test_basemaps.py new file mode 100644 index 0000000..5a066ca --- /dev/null +++ b/tests/test_basemaps.py @@ -0,0 +1,50 @@ +"""Test basemaps related API enpoints.""" +import pytest + +from datawrapper import Datawrapper + + +def test_get_basemaps(): + """Test the get_basemaps method.""" + dw = Datawrapper() + basemaps_list = dw.get_basemaps() + assert len(basemaps_list) > 0 + + +def test_get_basemap(): + """Test the get_basemap method.""" + dw = Datawrapper() + + # Test the standard query + basemap_info = dw.get_basemap("iowa-counties") + assert isinstance(basemap_info, dict) + assert basemap_info['meta']['slug'] == 'usa-iowa-counties' + assert basemap_info['meta']['projection'] == { + 'rotate': [93.49589653689938, -42.075128883839746], + 'type': 'geoAzimuthalEqualArea' + } + + # Test the projection kwarg + basemap_info = dw.get_basemap("iowa-counties", wgs84=True) + assert basemap_info['meta']['projection'] == { + 'type': 'geoAzimuthalEqualArea' + } + + # Test a missing slug + with pytest.raises(Exception): + dw.get_basemap("iowa-counties-zzz") + + +def test_get_basemap_key(): + """Test the get_basemap_key method.""" + dw = Datawrapper() + + # Test the standard query + basemap_key = dw.get_basemap_key("iowa-counties", 'GEOID') + assert isinstance(basemap_key, dict) + assert basemap_key['label'] == 'FIPS' + assert len(basemap_key['values']) == 99 + + # Test a missing slug + with pytest.raises(Exception): + dw.get_basemap_key("iowa-counties-zzz", 'GEOID') \ No newline at end of file