Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions docs/basic_functionality.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion docs/whats_new.md
Original file line number Diff line number Diff line change
@@ -1,13 +1,19 @@
# 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.


### 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.
Expand Down
2 changes: 1 addition & 1 deletion mkdocs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
123 changes: 123 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
6 changes: 4 additions & 2 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions wbdata/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__",
]
70 changes: 53 additions & 17 deletions wbdata/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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'
Expand All @@ -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)
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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'
Expand All @@ -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
Expand All @@ -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(
Expand Down