diff --git a/docs/basic_functionality.md b/docs/basic_functionality.md index f8db291..157e123 100644 --- a/docs/basic_functionality.md +++ b/docs/basic_functionality.md @@ -33,6 +33,20 @@ These functions require Pandas to be installed to work. show_root_full_path: false heading_level: 4 +### Excluding Aggregates + +By default, querying data for all countries (`country="all"`) includes regional and income aggregates (such as "World", "Arab World", and "East Asia & Pacific"). + +To exclude aggregates and retrieve data for individual countries and territories, use the `country="exclude_aggregates"` parameter value in the data retrieval functions, or `exclude_aggregates=True` when retrieving country lists: + +```python +# Query data excluding regional and income aggregates +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="exclude_aggregates") + +# Or when fetching country metadata: +countries = wbdata.get_countries(exclude_aggregates=True) +``` + ## Metadata Retrieval These functions, for the most part, are for finding the parameters you want to diff --git a/docs/whats_new.md b/docs/whats_new.md index 3dbd64a..9bad67d 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -1,5 +1,11 @@ # What's New +## Features in Development + +* Ability to exclude World Bank regional and group aggregates: + * `exclude_aggregates: bool = False` parameter on `get_countries`. + * Magic value `country="exclude_aggregates"` on data retrieval functions (`get_dataframe`, `get_series`, and `get_data`). + ## What's new in wbdata 1.0 The 1.0 release of `wbdata` is not *quite* a full rewrite, but is pretty much the next best thing. The architecture has been reworked, function and argument names have been changed to be more consistent and clear, and a few dependencies have been added for better and more reliable functionality. @@ -7,7 +13,7 @@ The 1.0 release of `wbdata` is not *quite* a full rewrite, but is pretty much th ### Features -* Date arguments can now be strings, not just `datetime.datetime` objects. Strings can be in the year, month, or quarter formats used by the World Bank API or in any other format that can be handled by [dateparser][https://dateparser.readthedocs.io/en/latest/]. +* Date arguments can now be strings, not just `datetime.datetime` objects. Strings can be in the year, month, or quarter formats used by the World Bank API or in any other format that can be handled by [dateparser](https://dateparser.readthedocs.io/en/latest/). * Default cache behavior can be configured with environment variables, including the path, TTL, and max number of items to cache. See [Cache Module documentation](reference/cache.md) for details. * Users can now create `Client` objects if they want to set cache behavior programmatically have multiple caches, or supply their own requests Session. * Caching is now provided using the [shelved_cache](https://github.com/mariushelf/shelved_cache) and [cachetools](https://github.com/tkem/cachetools/) libraries. Since a lot of annoying bugs seemed to come from wbdata's home-rolled cache implementation, this should be a good quality-of-life improvement for many people. diff --git a/mkdocs.yml b/mkdocs.yml index 5832bc0..520641a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -8,7 +8,7 @@ plugins: options: show_source: false members_order: source - docstrings_options: + docstring_options: returns_named_value: false returns_multiple_items: false watch: diff --git a/pyproject.toml b/pyproject.toml index 1df2e28..ebf821b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,8 @@ dependencies = [ pandas = ["pandas>=1,<3"] docs = [ "mkdocs>=1.5.3,<2", - "mkdocstrings[python]>=0.30,<1", + "mkdocstrings[python]>=0.30,<2", + "pymdown-extensions<12.2", ] [project.urls] diff --git a/tests/test_client.py b/tests/test_client.py index 9bb278a..6daa111 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -807,3 +807,126 @@ def test_get_dataframe(mock_client, results, indicators, keep_levels, expected): with mock.patch.object(mock_client, "get_series", mock.Mock(side_effect=results)): got = mock_client.get_dataframe(indicators=indicators, keep_levels=keep_levels) pd.testing.assert_frame_equal(got.loc[expected.index], expected) + + +def test_get_countries_exclude_aggregates(mock_client): + mock_data = [ + { + "id": "USA", + "iso2Code": "US", + "name": "United States", + "region": {"id": "NAC", "value": "North America"}, + }, + { + "id": "ABW", + "iso2Code": "AW", + "name": "Aruba", + "region": {"id": "LCN", "value": "Latin America & Caribbean "}, + }, + { + "id": "WLD", + "iso2Code": "1W", + "name": "World", + "region": {"id": "NA", "value": "Aggregates"}, + }, + ] + mock_client.fetcher.fetch = mock.Mock(return_value=mock_data) + + # Default: returns all 3 + assert len(mock_client.get_countries()) == 3 + + # Exclude aggregates: returns USA and Aruba + no_aggs = mock_client.get_countries(exclude_aggregates=True) + assert [c["id"] for c in no_aggs] == ["USA", "ABW"] + + # When country_id is passed + with mock.patch.object( + mock_client, + "_id_only_query", + mock.Mock(return_value=client.SearchResult([mock_data[2]])), + ): + res = mock_client.get_countries("WLD", exclude_aggregates=True) + assert len(res) == 0 + + +def test_get_data_exclude_aggregates(mock_client): + mock_countries = client.SearchResult( + [ + {"id": "USA", "name": "United States"}, + {"id": "ABW", "name": "Aruba"}, + ] + ) + with mock.patch.object( + mock_client, "get_countries", mock.Mock(return_value=mock_countries) + ) as mock_gc: + mock_client.fetcher.fetch = mock.Mock(return_value=[]) + mock_client.get_data("FOO", country="exclude_aggregates") + mock_gc.assert_called_once_with(exclude_aggregates=True, skip_cache=False) + mock_client.fetcher.fetch.assert_called_once_with( + url="https://api.worldbank.org/v2/countries/USA;ABW/indicators/FOO", + params={}, + skip_cache=False, + ) + + +def test_get_series_exclude_aggregates(mock_client): + mock_result = fetcher.Result( + [{"country": {"value": "United States"}, "date": "2020", "value": "100"}], + last_updated=dt.datetime(2023, 1, 1), + ) + with mock.patch.object( + mock_client, "get_data", mock.Mock(return_value=mock_result) + ) as mock_gd: + mock_client.get_series("FOO", country="exclude_aggregates") + mock_gd.assert_called_once_with( + indicator="FOO", + country="exclude_aggregates", + date=None, + freq="Y", + source=None, + parse_dates=False, + skip_cache=False, + ) + + +def test_get_dataframe_exclude_aggregates(mock_client): + mock_countries = client.SearchResult( + [ + {"id": "USA", "name": "United States"}, + {"id": "ABW", "name": "Aruba"}, + ] + ) + mock_series = client.Series( + [5.0, 6.0], + index=pd.MultiIndex.from_tuples( + tuples=( + ("usa", "2023"), + ("usa", "2024"), + ), + names=["country", "date"], + ), + ) + with ( + mock.patch.object( + mock_client, "get_countries", mock.Mock(return_value=mock_countries) + ) as mock_gc, + mock.patch.object( + mock_client, + "get_series", + mock.Mock(return_value=mock_series), + ) as mock_gs, + ): + mock_client.get_dataframe( + {"FOO": "foo"}, country="exclude_aggregates", skip_cache=True + ) + mock_gc.assert_called_once_with(exclude_aggregates=True, skip_cache=True) + mock_gs.assert_called_once_with( + indicator="FOO", + country=["USA", "ABW"], + date=None, + freq="Y", + source=None, + parse_dates=False, + keep_levels=True, + skip_cache=True, + ) diff --git a/uv.lock b/uv.lock index 9801e3c..677cf79 100644 --- a/uv.lock +++ b/uv.lock @@ -261,7 +261,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -1240,6 +1240,7 @@ dependencies = [ docs = [ { name = "mkdocs" }, { name = "mkdocstrings", extra = ["python"] }, + { name = "pymdown-extensions" }, ] pandas = [ { name = "pandas" }, @@ -1265,8 +1266,9 @@ requires-dist = [ { name = "dateparser", specifier = ">=1.2.0,<2" }, { name = "decorator", specifier = ">=5.1.1,<6" }, { name = "mkdocs", marker = "extra == 'docs'", specifier = ">=1.5.3,<2" }, - { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.30,<1" }, + { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.30,<2" }, { name = "pandas", marker = "extra == 'pandas'", specifier = ">=1,<3" }, + { name = "pymdown-extensions", marker = "extra == 'docs'", specifier = "<12.2" }, { name = "requests", specifier = ">=2.0,<3" }, { name = "shelved-cache", specifier = ">=0.3.1,<0.4" }, { name = "tabulate", specifier = ">=0.8.5,<1" }, diff --git a/wbdata/__init__.py b/wbdata/__init__.py index 323a2b1..036d980 100644 --- a/wbdata/__init__.py +++ b/wbdata/__init__.py @@ -25,3 +25,18 @@ def get_default_client() -> Client: get_lendingtypes = get_default_client().get_lendingtypes get_sources = get_default_client().get_sources get_topics = get_default_client().get_topics + +__all__ = [ + "Client", + "get_countries", + "get_data", + "get_dataframe", + "get_default_client", + "get_incomelevels", + "get_indicators", + "get_lendingtypes", + "get_series", + "get_sources", + "get_topics", + "__version__", +] diff --git a/wbdata/client.py b/wbdata/client.py index f64e38f..cdf7e4c 100644 --- a/wbdata/client.py +++ b/wbdata/client.py @@ -148,6 +148,16 @@ def _filter_by_pattern( return (row for row in rows if pattern.search(row["name"])) +def _is_aggregate_country(country: dict[str, Any]) -> bool: + """Return True if a country dictionary represents an aggregate region/group.""" + region = country.get("region") + if isinstance(region, dict): + return region.get("value") == "Aggregates" or region.get("id") == "NA" + if isinstance(region, str): + return region in ("Aggregates", "NA") + return False + + @dataclasses.dataclass class Client: """ @@ -197,7 +207,8 @@ def get_data( Parameters: indicator: the desired indicator code - country: a country code, sequence of country codes, or "all" (default) + country: a country code, sequence of country codes, "all" (default), + or "exclude_aggregates" date: the desired date as a string, datetime object or a 2-tuple with start and end dates freq: the desired periodicity of the data, one of 'Y' (yearly), 'M' @@ -212,6 +223,13 @@ def get_data( Returns: A list of dictionaries of observations """ + if isinstance(country, str) and country.strip().lower() == "exclude_aggregates": + country = [ + c["id"] + for c in self.get_countries( + exclude_aggregates=True, skip_cache=skip_cache + ) + ] url = COUNTRIES_URL try: c_part = _parse_value_or_iterable(country) @@ -322,12 +340,13 @@ def get_countries( query: str | re.Pattern | None = None, incomelevel: int | str | Sequence[int | str] | None = None, lendingtype: int | str | Sequence[int | str] | None = None, + exclude_aggregates: bool = False, skip_cache: bool = False, ) -> SearchResult: """ Retrieve information on one or more country or regional aggregates. - You can filter your results by specifying `query, `incomelevel`, or + You can filter your results by specifying `query`, `incomelevel`, or `lendingtype`. Specifying `query` will only return countries with names that match the query as a regular expression. If a string is supplied, the match is case insensitive. @@ -341,27 +360,36 @@ def get_countries( query: a regular expression on which to filter results incomelevel: desired incomelevel id or ids on which to filter results lendingtype: desired lendingtype id or ids on which to filter results + exclude_aggregates: if True, exclude regional and group aggregates + from results. Defaults to False. skip_cache: bypass cache when downloading Returns: list of dictionaries describing countries """ + raw_results: Iterable[dict[str, Any]] if country_id: if incomelevel or lendingtype or query: raise ValueError("Can't specify country_id and aggregates") - return self._id_only_query(COUNTRIES_URL, country_id, skip_cache=skip_cache) - params = {} - if incomelevel: - params["incomeLevel"] = _parse_value_or_iterable(incomelevel) - if lendingtype: - params["lendingType"] = _parse_value_or_iterable(lendingtype) - results = self.fetcher.fetch( - url=COUNTRIES_URL, params=params, skip_cache=skip_cache - ) - if query: - results = _filter_by_pattern(results, query) - return SearchResult(results) + raw_results = self._id_only_query( + COUNTRIES_URL, country_id, skip_cache=skip_cache + ) + else: + params = {} + if incomelevel: + params["incomeLevel"] = _parse_value_or_iterable(incomelevel) + if lendingtype: + params["lendingType"] = _parse_value_or_iterable(lendingtype) + fetched = self.fetcher.fetch( + url=COUNTRIES_URL, params=params, skip_cache=skip_cache + ) + raw_results = _filter_by_pattern(fetched, query) if query else fetched + + if exclude_aggregates: + raw_results = [r for r in raw_results if not _is_aggregate_country(r)] + + return SearchResult(raw_results) def get_indicators( self, @@ -436,7 +464,8 @@ def get_series( Parameters: indicator: the desired indicator code - country: a country code, sequence of country codes, or "all" (default) + country: a country code, sequence of country codes, "all" (default), + or "exclude_aggregates" date: the desired date as a string, datetime object or a 2-tuple with start and end dates freq: the desired periodicity of the data, one of 'Y' (yearly), 'M' @@ -507,7 +536,8 @@ def get_dataframe( Parameters: indicators: An dictionary where the keys are desired indicators and the values are the desired column names - country: a country code, sequence of country codes, or "all" (default) + country: a country code, sequence of country codes, "all" (default), + or "exclude_aggregates" date: the desired date as a string, datetime object or a 2-tuple with start and end dates freq: the desired periodicity of the data, one of 'Y' (yearly), 'M' @@ -517,7 +547,6 @@ def get_dataframe( to 2, World Development Indicators) parse_dates: if True, convert date field to a datetime.datetime object. - skip_cache: bypass the cache when downloading keep_levels: if True don't reduce the number of index levels returned if only getting one date or country skip_cache: bypass the cache when downloading @@ -533,6 +562,13 @@ def get_dataframe( levels only have one value, the country level is dropped. """ + if isinstance(country, str) and country.strip().lower() == "exclude_aggregates": + country = [ + c["id"] + for c in self.get_countries( + exclude_aggregates=True, skip_cache=skip_cache + ) + ] df = DataFrame( serieses={ name: self.get_series(