From be45c8bc58cbc520ca9651b9a50cfe5987dc3ee3 Mon Sep 17 00:00:00 2001 From: Nde Dilan Date: Sun, 11 Oct 2026 08:50:31 +0200 Subject: [PATCH 1/4] feat: allow excluding aggregates and non-sovereign countries (#31) --- docs/basic_functionality.md | 24 ++ docs/whats_new.md | 10 +- tests/test_client.py | 164 ++++++++++ wbdata/__init__.py | 23 +- wbdata/client.py | 585 ++++++++++++++++++++++++++++++++++-- 5 files changed, 779 insertions(+), 27 deletions(-) diff --git a/docs/basic_functionality.md b/docs/basic_functionality.md index f8db291..10cd52c 100644 --- a/docs/basic_functionality.md +++ b/docs/basic_functionality.md @@ -33,6 +33,30 @@ These functions require Pandas to be installed to work. show_root_full_path: false heading_level: 4 +### Excluding Aggregates and Non-Sovereign Countries + +By default, querying data for all countries (`country="all"`) includes regional and income aggregates (such as "World", "Arab World", and "High income") as well as non-sovereign territories/dependencies (such as Puerto Rico, Greenland, and Aruba). + +You can easily exclude aggregates or filter down to sovereign countries only across `get_dataframe`, `get_series`, `get_data`, and `get_countries`: + +```python +# Exclude aggregates (returns 217 countries and territories) +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, include_aggregates=False) +# or equivalently using special country keyword: +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="countries") + +# Exclude both aggregates and non-sovereign territories (returns 194 sovereign states) +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, sovereign_only=True) +# or equivalently: +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="sovereign") + +# Similarly when retrieving country metadata: +countries = wbdata.get_countries(include_aggregates=False) +sovereign_countries = wbdata.get_countries(sovereign_only=True) +``` + +The tuple of recognized non-sovereign countries is also directly exported as `wbdata.NON_SOVEREIGN_COUNTRIES` (and `wbdata.non_sovereign_countries`). + ## 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..27ca41d 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -1,5 +1,13 @@ # What's New +## Features in Development + +* Ability to exclude World Bank aggregates and/or non-sovereign territories: + * `include_aggregates: bool = True` parameter on `get_dataframe`, `get_series`, `get_data`, and `get_countries`. + * `include_non_sovereign: bool = True` and `sovereign_only: bool = False` parameters to filter non-sovereign countries. + * Special string values for `country`: `country="countries"` (excluding aggregates) and `country="sovereign"` (sovereign countries only). + * `wbdata.NON_SOVEREIGN_COUNTRIES` list of non-sovereign territories exported at package level. + ## 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 +15,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/tests/test_client.py b/tests/test_client.py index 9bb278a..5c983a2 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -807,3 +807,167 @@ 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_non_sovereign_constants(): + import wbdata + + assert len(client.NON_SOVEREIGN_COUNTRIES) == 23 + assert "Puerto Rico (US)" in client.NON_SOVEREIGN_COUNTRIES + assert "Aruba" in client.NON_SOVEREIGN_COUNTRIES + assert client.non_sovereign_countries == list(client.NON_SOVEREIGN_COUNTRIES) + assert wbdata.NON_SOVEREIGN_COUNTRIES == client.NON_SOVEREIGN_COUNTRIES + assert wbdata.non_sovereign_countries == client.non_sovereign_countries + + +def test_get_countries_filtering(mock_client): + mock_data = [ + { + "id": "USA", + "iso2Code": "US", + "name": "United States", + "incomeLevel": {"id": "HIC", "value": "High income"}, + "region": {"id": "NAC", "value": "North America"}, + }, + { + "id": "ABW", + "iso2Code": "AW", + "name": "Aruba", + "incomeLevel": {"id": "HIC", "value": "High income"}, + "region": {"id": "LCN", "value": "Latin America & Caribbean "}, + }, + { + "id": "WLD", + "iso2Code": "1W", + "name": "World", + "incomeLevel": {"id": "NA", "value": "Aggregates"}, + "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(include_aggregates=False) + assert [c["id"] for c in no_aggs] == ["USA", "ABW"] + + # exclude_aggregates alias + no_aggs_alias = mock_client.get_countries(exclude_aggregates=True) + assert [c["id"] for c in no_aggs_alias] == ["USA", "ABW"] + + # Exclude non-sovereign: returns USA and World + no_ns = mock_client.get_countries(include_non_sovereign=False) + assert [c["id"] for c in no_ns] == ["USA", "WLD"] + + # Sovereign only: returns USA only + sov = mock_client.get_countries(sovereign_only=True) + assert [c["id"] for c in sov] == ["USA"] + + +def test_get_data_filtering(mock_client): + mock_rows = fetcher.Result( + [ + { + "country": {"id": "US", "value": "United States"}, + "countryiso3code": "USA", + "date": "2020", + "value": 100, + }, + { + "country": {"id": "AW", "value": "Aruba"}, + "countryiso3code": "ABW", + "date": "2020", + "value": 200, + }, + { + "country": {"id": "1W", "value": "World"}, + "countryiso3code": "WLD", + "date": "2020", + "value": 300, + }, + ], + last_updated=dt.datetime(2023, 1, 1), + ) + mock_client.fetcher.fetch = mock.Mock(return_value=mock_rows) + + # Default: returns all 3 + res_all = mock_client.get_data("FOO") + assert len(res_all) == 3 + + # Exclude aggregates + res_no_aggs = mock_client.get_data("FOO", include_aggregates=False) + assert [r["countryiso3code"] for r in res_no_aggs] == ["USA", "ABW"] + assert res_no_aggs.last_updated == dt.datetime(2023, 1, 1) + + # Exclude non-sovereign + res_no_ns = mock_client.get_data("FOO", include_non_sovereign=False) + assert [r["countryiso3code"] for r in res_no_ns] == ["USA", "WLD"] + + # Sovereign only + res_sov = mock_client.get_data("FOO", sovereign_only=True) + assert [r["countryiso3code"] for r in res_sov] == ["USA"] + + # Special country strings + res_countries = mock_client.get_data("FOO", country="countries") + assert [r["countryiso3code"] for r in res_countries] == ["USA", "ABW"] + + res_countries_only = mock_client.get_data("FOO", country="countries_only") + assert [r["countryiso3code"] for r in res_countries_only] == ["USA", "ABW"] + + res_sov_country = mock_client.get_data("FOO", country="sovereign") + assert [r["countryiso3code"] for r in res_sov_country] == ["USA"] + + +def test_get_series_filtering_args(mock_client): + with mock.patch.object(mock_client, "get_data", mock.Mock(return_value=[])): + mock_client.get_series( + "FOO", + include_aggregates=False, + include_non_sovereign=False, + sovereign_only=True, + exclude_aggregates=True, + ) + mock_client.get_data.assert_called_once_with( + indicator="FOO", + country="all", + date=None, + freq="Y", + source=None, + parse_dates=False, + skip_cache=False, + include_aggregates=False, + include_non_sovereign=False, + sovereign_only=True, + exclude_aggregates=True, + ) + + +def test_get_dataframe_filtering_args(mock_client): + with mock.patch.object( + mock_client, + "get_series", + mock.Mock(return_value=client.Series([1.0], index=pd.Index(["USA"]))), + ): + mock_client.get_dataframe( + {"FOO": "foo"}, + include_aggregates=False, + include_non_sovereign=False, + sovereign_only=True, + exclude_aggregates=True, + ) + mock_client.get_series.assert_called_once_with( + indicator="FOO", + country="all", + date=None, + freq="Y", + source=None, + parse_dates=False, + keep_levels=True, + skip_cache=False, + include_aggregates=False, + include_non_sovereign=False, + sovereign_only=True, + exclude_aggregates=True, + ) diff --git a/wbdata/__init__.py b/wbdata/__init__.py index 323a2b1..16343fd 100644 --- a/wbdata/__init__.py +++ b/wbdata/__init__.py @@ -4,7 +4,11 @@ from functools import lru_cache -from .client import Client +from .client import ( + NON_SOVEREIGN_COUNTRIES, + Client, + non_sovereign_countries, +) from .version import __version__ @@ -25,3 +29,20 @@ 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__ = [ + "NON_SOVEREIGN_COUNTRIES", + "non_sovereign_countries", + "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..34761eb 100644 --- a/wbdata/client.py +++ b/wbdata/client.py @@ -30,6 +30,341 @@ SOURCE_URL = f"{BASE_URL}/sources" TOPIC_URL = f"{BASE_URL}/topics" +NON_SOVEREIGN_COUNTRIES: tuple[str, ...] = ( + "American Samoa", + "Aruba", + "Bermuda", + "British Virgin Islands", + "Cayman Islands", + "Channel Islands", + "Curacao", + "Faroe Islands", + "French Polynesia", + "Gibraltar", + "Greenland", + "Guam", + "Hong Kong SAR, China", + "Isle of Man", + "Macao SAR, China", + "New Caledonia", + "Northern Mariana Islands", + "Puerto Rico (US)", + "Sint Maarten (Dutch part)", + "St. Martin (French part)", + "Turks and Caicos Islands", + "Virgin Islands (U.S.)", + "West Bank and Gaza", +) +non_sovereign_countries = list(NON_SOVEREIGN_COUNTRIES) + +NON_SOVEREIGN_IDS: frozenset[str] = frozenset( + { + "ABW", + "ASM", + "BMU", + "CHI", + "CUW", + "CYM", + "FRO", + "GIB", + "GRL", + "GUM", + "HKG", + "IMN", + "MAC", + "MAF", + "MNP", + "NCL", + "PRI", + "PSE", + "PYF", + "SXM", + "TCA", + "VGB", + "VIR", + } +) + +NON_SOVEREIGN_ISO2: frozenset[str] = frozenset( + { + "AW", + "AS", + "BM", + "JG", + "CW", + "KY", + "FO", + "GI", + "GL", + "GU", + "HK", + "IM", + "MO", + "MF", + "MP", + "NC", + "PR", + "PS", + "PF", + "SX", + "TC", + "VG", + "VI", + } +) + +AGGREGATE_IDS: frozenset[str] = frozenset( + { + "AFE", + "AFR", + "AFW", + "ARB", + "BEA", + "BEC", + "BHI", + "BLA", + "BMN", + "BSS", + "CAA", + "CEA", + "CEB", + "CEU", + "CLA", + "CME", + "CSA", + "CSS", + "DEA", + "DEC", + "DLA", + "DMN", + "DNS", + "DSA", + "DSF", + "DSS", + "EAP", + "EAR", + "EAS", + "ECA", + "ECS", + "EMU", + "EUU", + "FXS", + "HIC", + "HPC", + "IBB", + "IBD", + "IBT", + "IDA", + "IDB", + "IDX", + "INX", + "LAC", + "LCN", + "LDC", + "LIC", + "LMC", + "LMY", + "LTE", + "MDE", + "MEA", + "MIC", + "MNA", + "NAC", + "NAF", + "NRS", + "NXS", + "OED", + "OSS", + "PRE", + "PSS", + "PST", + "RRS", + "SAS", + "SSA", + "SSF", + "SST", + "SXZ", + "TEA", + "TEC", + "TLA", + "TMN", + "TSA", + "TSS", + "UMC", + "WLD", + "XZN", + } +) + +AGGREGATE_ISO2: frozenset[str] = frozenset( + { + "1A", + "1W", + "4E", + "7E", + "8S", + "B1", + "B2", + "B3", + "B4", + "B6", + "B7", + "B8", + "C4", + "C5", + "C6", + "C7", + "C8", + "C9", + "D1", + "D2", + "D3", + "D4", + "D5", + "D6", + "D7", + "EU", + "F1", + "F6", + "M1", + "M2", + "N6", + "OE", + "S1", + "S2", + "S3", + "S4", + "T2", + "T3", + "T4", + "T5", + "T6", + "T7", + "V1", + "V2", + "V3", + "V4", + "XC", + "XD", + "XE", + "XF", + "XG", + "XH", + "XI", + "XJ", + "XL", + "XM", + "XN", + "XO", + "XP", + "XQ", + "XR", + "XS", + "XT", + "XU", + "XY", + "XZ", + "Z4", + "Z7", + "ZB", + "ZF", + "ZG", + "ZH", + "ZI", + "ZJ", + "ZQ", + "ZT", + } +) + +AGGREGATE_NAMES: frozenset[str] = frozenset( + { + "Africa", + "Africa Eastern and Southern", + "Africa Western and Central", + "Arab World", + "Caribbean small states", + "Central Europe and the Baltics", + "Early-demographic dividend", + "East Asia & Pacific", + "East Asia & Pacific (IBRD-only countries)", + "East Asia & Pacific (IDA & IBRD countries)", + "East Asia & Pacific (IDA-eligible countries)", + "East Asia & Pacific (excluding high income)", + "East Asia and the Pacific (IFC classification)", + "Euro area", + "Europe & Central Asia", + "Europe & Central Asia (IBRD-only countries)", + "Europe & Central Asia (IDA & IBRD countries)", + "Europe & Central Asia (IDA-eligible countries)", + "Europe & Central Asia (excluding high income)", + "Europe and Central Asia (IFC classification)", + "European Union", + "Heavily indebted poor countries (HIPC)", + "High income", + "IBRD countries classified as high income", + "IBRD only", + "IBRD, including blend", + "IDA & IBRD total", + ( + "IDA countries classified as fragile situations, " + "excluding Sub-Saharan Africa" + ), + "IDA countries in Sub-Saharan Africa classified as fragile situations ", + "IDA countries in Sub-Saharan Africa not classified as fragile situations ", + ( + "IDA countries not classified as fragile situations, " + "excluding Sub-Saharan Africa" + ), + "IDA only", + "IDA total", + "Late-demographic dividend", + "Latin America & Caribbean ", + "Latin America & Caribbean (excluding high income)", + "Latin America & the Caribbean (IBRD-only countries)", + "Latin America & the Caribbean (IDA & IBRD countries)", + "Latin America & the Caribbean (IDA-eligible countries)", + "Latin America and the Caribbean (IFC classification)", + "Least developed countries: UN classification", + "Low & middle income", + "Low income", + "Lower middle income", + "Middle East (developing only)", + "Middle East and North Africa (IFC classification)", + "Middle East, North Africa, Afghanistan & Pakistan", + "Middle East, North Africa, Afghanistan & Pakistan (IBRD only)", + "Middle East, North Africa, Afghanistan & Pakistan (IDA & IBRD)", + "Middle East, North Africa, Afghanistan & Pakistan (IDA total)", + "Middle East, North Africa, Afghanistan & Pakistan (excluding high income)", + "Middle income", + "Non-resource rich Sub-Saharan Africa countries", + "North Africa", + "North America", + "Not classified", + "OECD members", + "Other small states", + "Pacific island small states", + "Post-demographic dividend", + "Pre-demographic dividend", + "Resource rich Sub-Saharan Africa countries", + "Small states", + "South Asia", + "South Asia (IDA & IBRD)", + "South Asia (IDA-eligible countries)", + "South Asia (IFC classification)", + "Sub-Saharan Africa ", + "Sub-Saharan Africa (IBRD-only countries)", + "Sub-Saharan Africa (IDA & IBRD countries)", + "Sub-Saharan Africa (IDA-eligible countries)", + "Sub-Saharan Africa (IFC classification)", + "Sub-Saharan Africa (excluding high income)", + "Sub-Saharan Africa excluding South Africa", + "Sub-Saharan Africa excluding South Africa and Nigeria", + "Upper middle income", + "World", + } +) + class SearchResult(list): """ @@ -148,6 +483,91 @@ def _filter_by_pattern( return (row for row in rows if pattern.search(row["name"])) +def _is_aggregate_entity(entity: dict[str, Any]) -> bool: + """Return True if an entity dict represents an aggregate region or group.""" + income_level = entity.get("incomeLevel") + if isinstance(income_level, dict) and income_level.get("value") == "Aggregates": + return True + region = entity.get("region") + if isinstance(region, dict) and ( + region.get("value") == "Aggregates" or region.get("id") == "NA" + ): + return True + lending_type = entity.get("lendingType") + if isinstance(lending_type, dict) and lending_type.get("value") == "Aggregates": + return True + if entity.get("id") in AGGREGATE_IDS: + return True + if entity.get("iso2Code") in AGGREGATE_ISO2: + return True + return entity.get("name") in AGGREGATE_NAMES + + +def _is_non_sovereign_entity(entity: dict[str, Any]) -> bool: + """Return True if an entity dict represents a non-sovereign country/territory.""" + if entity.get("name") in NON_SOVEREIGN_COUNTRIES: + return True + if entity.get("id") in NON_SOVEREIGN_IDS: + return True + return entity.get("iso2Code") in NON_SOVEREIGN_ISO2 + + +def _is_aggregate_row(row: dict[str, Any]) -> bool: + """Return True if an observation row represents an aggregate.""" + iso3 = row.get("countryiso3code") + if iso3 and iso3 in AGGREGATE_IDS: + return True + country_obj = row.get("country") + if isinstance(country_obj, dict): + if country_obj.get("id") in AGGREGATE_ISO2: + return True + if country_obj.get("value") in AGGREGATE_NAMES: + return True + return False + + +def _is_non_sovereign_row(row: dict[str, Any]) -> bool: + """Return True if an observation row represents a non-sovereign country.""" + iso3 = row.get("countryiso3code") + if iso3 and iso3 in NON_SOVEREIGN_IDS: + return True + country_obj = row.get("country") + if isinstance(country_obj, dict): + if country_obj.get("value") in NON_SOVEREIGN_COUNTRIES: + return True + if country_obj.get("id") in NON_SOVEREIGN_ISO2: + return True + return False + + +def _resolve_country_and_flags( + country: str | Sequence[str], + include_aggregates: bool, + include_non_sovereign: bool, + sovereign_only: bool, + exclude_aggregates: bool | None, +) -> tuple[str | Sequence[str], bool, bool]: + """Resolve and normalize country and filter flags.""" + if isinstance(country, str): + c_lower = country.strip().lower() + if c_lower in ("countries", "countries_only"): + country = "all" + include_aggregates = False + elif c_lower == "sovereign": + country = "all" + include_aggregates = False + include_non_sovereign = False + + if exclude_aggregates is not None: + include_aggregates = not exclude_aggregates + + if sovereign_only: + include_aggregates = False + include_non_sovereign = False + + return country, include_aggregates, include_non_sovereign + + @dataclasses.dataclass class Client: """ @@ -191,13 +611,19 @@ def get_data( source: int | str | Sequence[int | str] | None = None, parse_dates: bool = False, skip_cache: bool = False, + include_aggregates: bool = True, + include_non_sovereign: bool = True, + sovereign_only: bool = False, + exclude_aggregates: bool | None = None, ) -> fetcher.Result: """ Retrieve indicators for given countries and years 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, or "all" (default). + Also accepts special strings "countries" / "countries_only" (exclude + aggregates) and "sovereign" (sovereign countries only). 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' @@ -208,10 +634,26 @@ def get_data( parse_dates: if True, convert date field to a datetime.datetime object. skip_cache: bypass the cache when downloading + include_aggregates: if False, exclude aggregate regions and groups + from the returned data. Defaults to True. + include_non_sovereign: if False, exclude non-sovereign + countries/territories from the returned data. Defaults to True. + sovereign_only: if True, exclude both aggregates and non-sovereign + countries. Defaults to False. + exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: A list of dictionaries of observations """ + country, include_aggregates, include_non_sovereign = ( + _resolve_country_and_flags( + country=country, + include_aggregates=include_aggregates, + include_non_sovereign=include_non_sovereign, + sovereign_only=sovereign_only, + exclude_aggregates=exclude_aggregates, + ) + ) url = COUNTRIES_URL try: c_part = _parse_value_or_iterable(country) @@ -226,6 +668,19 @@ def get_data( data = self.fetcher.fetch(url=url, params=params, skip_cache=skip_cache) if parse_dates: dates.parse_row_dates(data) + if (not include_aggregates or not include_non_sovereign) and isinstance( + data, list + ): + filtered_rows = [ + row + for row in data + if (include_aggregates or not _is_aggregate_row(row)) + and (include_non_sovereign or not _is_non_sovereign_row(row)) + ] + data = fetcher.Result( + filtered_rows, + last_updated=getattr(data, "last_updated", None), + ) return data def _id_only_query(self, url: str, id_: Any, skip_cache: bool) -> SearchResult: @@ -323,11 +778,15 @@ def get_countries( incomelevel: int | str | Sequence[int | str] | None = None, lendingtype: int | str | Sequence[int | str] | None = None, skip_cache: bool = False, + include_aggregates: bool = True, + include_non_sovereign: bool = True, + sovereign_only: bool = False, + exclude_aggregates: bool | None = None, ) -> 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. @@ -342,26 +801,51 @@ def get_countries( incomelevel: desired incomelevel id or ids on which to filter results lendingtype: desired lendingtype id or ids on which to filter results skip_cache: bypass cache when downloading + include_aggregates: if False, exclude aggregate regions and groups + from results. Defaults to True. + include_non_sovereign: if False, exclude non-sovereign + countries/territories from results. Defaults to True. + sovereign_only: if True, exclude both aggregates and non-sovereign + countries. Defaults to False. + exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: list of dictionaries describing countries """ + if exclude_aggregates is not None: + include_aggregates = not exclude_aggregates + if sovereign_only: + include_aggregates = False + include_non_sovereign = False + + 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 not include_aggregates or not include_non_sovereign: + raw_results = [ + row + for row in raw_results + if (include_aggregates or not _is_aggregate_entity(row)) + and (include_non_sovereign or not _is_non_sovereign_entity(row)) + ] + + return SearchResult(raw_results) def get_indicators( self, @@ -428,6 +912,10 @@ def get_series( name: str = "value", keep_levels: bool = False, skip_cache: bool = False, + include_aggregates: bool = True, + include_non_sovereign: bool = True, + sovereign_only: bool = False, + exclude_aggregates: bool | None = None, ) -> Series: """ Retrieve data for a single indicator as a pandas Series. @@ -436,7 +924,9 @@ 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, or "all" (default). + Also accepts special strings "countries" / "countries_only" (exclude + aggregates) and "sovereign" (sovereign countries only). 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' @@ -446,11 +936,17 @@ def get_series( to 2, World Development Indicators) parse_dates: if True, convert date field to a datetime.datetime object. - skip_cache: bypass the cache when downloading name: the desired name for the pandas Series 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 + include_aggregates: if False, exclude aggregate regions and groups + from the returned data. Defaults to True. + include_non_sovereign: if False, exclude non-sovereign + countries/territories from the returned data. Defaults to True. + sovereign_only: if True, exclude both aggregates and non-sovereign + countries. Defaults to False. + exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: Series with the requested data. The index of the series depends on @@ -462,6 +958,16 @@ def get_series( will be dropped. If `keep_levels` is `False` and both levels only have one value, the country level is dropped. """ + extra_kwargs: dict[str, Any] = {} + if not include_aggregates: + extra_kwargs["include_aggregates"] = include_aggregates + if not include_non_sovereign: + extra_kwargs["include_non_sovereign"] = include_non_sovereign + if sovereign_only: + extra_kwargs["sovereign_only"] = sovereign_only + if exclude_aggregates is not None: + extra_kwargs["exclude_aggregates"] = exclude_aggregates + raw_data = self.get_data( indicator=indicator, country=country, @@ -470,6 +976,7 @@ def get_series( source=source, parse_dates=parse_dates, skip_cache=skip_cache, + **extra_kwargs, ) df = pd.DataFrame( [[i["country"]["value"], i["date"], i["value"]] for i in raw_data], @@ -482,7 +989,7 @@ def get_series( df = df.set_index("country") else: df = df.set_index(["country", "date"]) - return Series(df[name], last_updated=raw_data.last_updated) + return Series(df[name], last_updated=getattr(raw_data, "last_updated", None)) @needs_pandas def get_dataframe( @@ -498,16 +1005,22 @@ def get_dataframe( parse_dates: bool = False, keep_levels: bool = False, skip_cache: bool = False, + include_aggregates: bool = True, + include_non_sovereign: bool = True, + sovereign_only: bool = False, + exclude_aggregates: bool | None = None, ) -> DataFrame: """ - Download a set of indicators and merge them into a pandas DataFrame. + Download a set of indicators and merge them into a pandas DataFrame. If pandas is not installed, a RuntimeError will be raised. Parameters: - indicators: An dictionary where the keys are desired indicators and the + indicators: A 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, or "all" (default). + Also accepts special strings "countries" / "countries_only" (exclude + aggregates) and "sovereign" (sovereign countries only). 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,10 +1030,16 @@ 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 + include_aggregates: if False, exclude aggregate regions and groups + from the returned data. Defaults to True. + include_non_sovereign: if False, exclude non-sovereign + countries/territories from the returned data. Defaults to True. + sovereign_only: if True, exclude both aggregates and non-sovereign + countries. Defaults to False. + exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: DataFrame with one column per indicator. The index of the DataFrame @@ -533,6 +1052,16 @@ def get_dataframe( levels only have one value, the country level is dropped. """ + extra_kwargs: dict[str, Any] = {} + if not include_aggregates: + extra_kwargs["include_aggregates"] = include_aggregates + if not include_non_sovereign: + extra_kwargs["include_non_sovereign"] = include_non_sovereign + if sovereign_only: + extra_kwargs["sovereign_only"] = sovereign_only + if exclude_aggregates is not None: + extra_kwargs["exclude_aggregates"] = exclude_aggregates + df = DataFrame( serieses={ name: self.get_series( @@ -544,12 +1073,18 @@ def get_dataframe( parse_dates=parse_dates, keep_levels=True, skip_cache=skip_cache, + **extra_kwargs, ) for indicator, name in indicators.items() } ) - if not keep_levels and len(set(df.index.get_level_values(0))) == 1: - df.index = df.index.droplevel(0) - elif not keep_levels and len(set(df.index.get_level_values(1))) == 1: - df.index = df.index.droplevel(1) + if ( + not keep_levels + and isinstance(df.index, pd.MultiIndex) + and len(df.index) > 0 + ): + if len(set(df.index.get_level_values(0))) == 1: + df.index = df.index.droplevel(0) + elif len(set(df.index.get_level_values(1))) == 1: + df.index = df.index.droplevel(1) return df From fe4773033b6087784d48ce749b136523ad7ea11f Mon Sep 17 00:00:00 2001 From: Nde Dilan Date: Sun, 11 Oct 2026 09:10:22 +0200 Subject: [PATCH 2/4] fix: correct docstring_options key in mkdocs.yml --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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: From c539e53ee06f760c221bfa82b5cd25a608d93aa6 Mon Sep 17 00:00:00 2001 From: Nde Dilan Date: Sun, 11 Oct 2026 09:29:18 +0200 Subject: [PATCH 3/4] fix(docs): pin pymdown-extensions<12.2 and allow mkdocstrings<2 --- pyproject.toml | 3 ++- uv.lock | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) 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/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" }, From 01fe8dbc8727b22adcab137269f9abd2a97894f2 Mon Sep 17 00:00:00 2001 From: Nde Dilan Date: Sat, 31 Oct 2026 01:34:47 +0200 Subject: [PATCH 4/4] refactor: simplify aggregate exclusion via country region and magic value (#31) --- docs/basic_functionality.md | 26 +- docs/whats_new.md | 8 +- tests/test_client.py | 177 +++++------ wbdata/__init__.py | 8 +- wbdata/client.py | 579 +++--------------------------------- 5 files changed, 120 insertions(+), 678 deletions(-) diff --git a/docs/basic_functionality.md b/docs/basic_functionality.md index 10cd52c..157e123 100644 --- a/docs/basic_functionality.md +++ b/docs/basic_functionality.md @@ -33,29 +33,19 @@ These functions require Pandas to be installed to work. show_root_full_path: false heading_level: 4 -### Excluding Aggregates and Non-Sovereign Countries +### Excluding Aggregates -By default, querying data for all countries (`country="all"`) includes regional and income aggregates (such as "World", "Arab World", and "High income") as well as non-sovereign territories/dependencies (such as Puerto Rico, Greenland, and Aruba). +By default, querying data for all countries (`country="all"`) includes regional and income aggregates (such as "World", "Arab World", and "East Asia & Pacific"). -You can easily exclude aggregates or filter down to sovereign countries only across `get_dataframe`, `get_series`, `get_data`, and `get_countries`: +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 -# Exclude aggregates (returns 217 countries and territories) -df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, include_aggregates=False) -# or equivalently using special country keyword: -df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="countries") - -# Exclude both aggregates and non-sovereign territories (returns 194 sovereign states) -df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, sovereign_only=True) -# or equivalently: -df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="sovereign") - -# Similarly when retrieving country metadata: -countries = wbdata.get_countries(include_aggregates=False) -sovereign_countries = wbdata.get_countries(sovereign_only=True) -``` +# Query data excluding regional and income aggregates +df = wbdata.get_dataframe({'SP.POP.TOTL': 'pop'}, country="exclude_aggregates") -The tuple of recognized non-sovereign countries is also directly exported as `wbdata.NON_SOVEREIGN_COUNTRIES` (and `wbdata.non_sovereign_countries`). +# Or when fetching country metadata: +countries = wbdata.get_countries(exclude_aggregates=True) +``` ## Metadata Retrieval diff --git a/docs/whats_new.md b/docs/whats_new.md index 27ca41d..9bad67d 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -2,11 +2,9 @@ ## Features in Development -* Ability to exclude World Bank aggregates and/or non-sovereign territories: - * `include_aggregates: bool = True` parameter on `get_dataframe`, `get_series`, `get_data`, and `get_countries`. - * `include_non_sovereign: bool = True` and `sovereign_only: bool = False` parameters to filter non-sovereign countries. - * Special string values for `country`: `country="countries"` (excluding aggregates) and `country="sovereign"` (sovereign countries only). - * `wbdata.NON_SOVEREIGN_COUNTRIES` list of non-sovereign territories exported at package level. +* 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 diff --git a/tests/test_client.py b/tests/test_client.py index 5c983a2..6daa111 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -809,38 +809,24 @@ def test_get_dataframe(mock_client, results, indicators, keep_levels, expected): pd.testing.assert_frame_equal(got.loc[expected.index], expected) -def test_non_sovereign_constants(): - import wbdata - - assert len(client.NON_SOVEREIGN_COUNTRIES) == 23 - assert "Puerto Rico (US)" in client.NON_SOVEREIGN_COUNTRIES - assert "Aruba" in client.NON_SOVEREIGN_COUNTRIES - assert client.non_sovereign_countries == list(client.NON_SOVEREIGN_COUNTRIES) - assert wbdata.NON_SOVEREIGN_COUNTRIES == client.NON_SOVEREIGN_COUNTRIES - assert wbdata.non_sovereign_countries == client.non_sovereign_countries - - -def test_get_countries_filtering(mock_client): +def test_get_countries_exclude_aggregates(mock_client): mock_data = [ { "id": "USA", "iso2Code": "US", "name": "United States", - "incomeLevel": {"id": "HIC", "value": "High income"}, "region": {"id": "NAC", "value": "North America"}, }, { "id": "ABW", "iso2Code": "AW", "name": "Aruba", - "incomeLevel": {"id": "HIC", "value": "High income"}, "region": {"id": "LCN", "value": "Latin America & Caribbean "}, }, { "id": "WLD", "iso2Code": "1W", "name": "World", - "incomeLevel": {"id": "NA", "value": "Aggregates"}, "region": {"id": "NA", "value": "Aggregates"}, }, ] @@ -850,124 +836,97 @@ def test_get_countries_filtering(mock_client): assert len(mock_client.get_countries()) == 3 # Exclude aggregates: returns USA and Aruba - no_aggs = mock_client.get_countries(include_aggregates=False) + no_aggs = mock_client.get_countries(exclude_aggregates=True) assert [c["id"] for c in no_aggs] == ["USA", "ABW"] - # exclude_aggregates alias - no_aggs_alias = mock_client.get_countries(exclude_aggregates=True) - assert [c["id"] for c in no_aggs_alias] == ["USA", "ABW"] - - # Exclude non-sovereign: returns USA and World - no_ns = mock_client.get_countries(include_non_sovereign=False) - assert [c["id"] for c in no_ns] == ["USA", "WLD"] - - # Sovereign only: returns USA only - sov = mock_client.get_countries(sovereign_only=True) - assert [c["id"] for c in sov] == ["USA"] + # 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_filtering(mock_client): - mock_rows = fetcher.Result( +def test_get_data_exclude_aggregates(mock_client): + mock_countries = client.SearchResult( [ - { - "country": {"id": "US", "value": "United States"}, - "countryiso3code": "USA", - "date": "2020", - "value": 100, - }, - { - "country": {"id": "AW", "value": "Aruba"}, - "countryiso3code": "ABW", - "date": "2020", - "value": 200, - }, - { - "country": {"id": "1W", "value": "World"}, - "countryiso3code": "WLD", - "date": "2020", - "value": 300, - }, - ], - last_updated=dt.datetime(2023, 1, 1), + {"id": "USA", "name": "United States"}, + {"id": "ABW", "name": "Aruba"}, + ] ) - mock_client.fetcher.fetch = mock.Mock(return_value=mock_rows) - - # Default: returns all 3 - res_all = mock_client.get_data("FOO") - assert len(res_all) == 3 - - # Exclude aggregates - res_no_aggs = mock_client.get_data("FOO", include_aggregates=False) - assert [r["countryiso3code"] for r in res_no_aggs] == ["USA", "ABW"] - assert res_no_aggs.last_updated == dt.datetime(2023, 1, 1) - - # Exclude non-sovereign - res_no_ns = mock_client.get_data("FOO", include_non_sovereign=False) - assert [r["countryiso3code"] for r in res_no_ns] == ["USA", "WLD"] - - # Sovereign only - res_sov = mock_client.get_data("FOO", sovereign_only=True) - assert [r["countryiso3code"] for r in res_sov] == ["USA"] - - # Special country strings - res_countries = mock_client.get_data("FOO", country="countries") - assert [r["countryiso3code"] for r in res_countries] == ["USA", "ABW"] - - res_countries_only = mock_client.get_data("FOO", country="countries_only") - assert [r["countryiso3code"] for r in res_countries_only] == ["USA", "ABW"] - - res_sov_country = mock_client.get_data("FOO", country="sovereign") - assert [r["countryiso3code"] for r in res_sov_country] == ["USA"] + 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_filtering_args(mock_client): - with mock.patch.object(mock_client, "get_data", mock.Mock(return_value=[])): - mock_client.get_series( - "FOO", - include_aggregates=False, - include_non_sovereign=False, - sovereign_only=True, - exclude_aggregates=True, - ) - mock_client.get_data.assert_called_once_with( +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="all", + country="exclude_aggregates", date=None, freq="Y", source=None, parse_dates=False, skip_cache=False, - include_aggregates=False, - include_non_sovereign=False, - sovereign_only=True, - exclude_aggregates=True, ) -def test_get_dataframe_filtering_args(mock_client): - with mock.patch.object( - mock_client, - "get_series", - mock.Mock(return_value=client.Series([1.0], index=pd.Index(["USA"]))), +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"}, - include_aggregates=False, - include_non_sovereign=False, - sovereign_only=True, - exclude_aggregates=True, + {"FOO": "foo"}, country="exclude_aggregates", skip_cache=True ) - mock_client.get_series.assert_called_once_with( + mock_gc.assert_called_once_with(exclude_aggregates=True, skip_cache=True) + mock_gs.assert_called_once_with( indicator="FOO", - country="all", + country=["USA", "ABW"], date=None, freq="Y", source=None, parse_dates=False, keep_levels=True, - skip_cache=False, - include_aggregates=False, - include_non_sovereign=False, - sovereign_only=True, - exclude_aggregates=True, + skip_cache=True, ) diff --git a/wbdata/__init__.py b/wbdata/__init__.py index 16343fd..036d980 100644 --- a/wbdata/__init__.py +++ b/wbdata/__init__.py @@ -4,11 +4,7 @@ from functools import lru_cache -from .client import ( - NON_SOVEREIGN_COUNTRIES, - Client, - non_sovereign_countries, -) +from .client import Client from .version import __version__ @@ -31,8 +27,6 @@ def get_default_client() -> Client: get_topics = get_default_client().get_topics __all__ = [ - "NON_SOVEREIGN_COUNTRIES", - "non_sovereign_countries", "Client", "get_countries", "get_data", diff --git a/wbdata/client.py b/wbdata/client.py index 34761eb..cdf7e4c 100644 --- a/wbdata/client.py +++ b/wbdata/client.py @@ -30,341 +30,6 @@ SOURCE_URL = f"{BASE_URL}/sources" TOPIC_URL = f"{BASE_URL}/topics" -NON_SOVEREIGN_COUNTRIES: tuple[str, ...] = ( - "American Samoa", - "Aruba", - "Bermuda", - "British Virgin Islands", - "Cayman Islands", - "Channel Islands", - "Curacao", - "Faroe Islands", - "French Polynesia", - "Gibraltar", - "Greenland", - "Guam", - "Hong Kong SAR, China", - "Isle of Man", - "Macao SAR, China", - "New Caledonia", - "Northern Mariana Islands", - "Puerto Rico (US)", - "Sint Maarten (Dutch part)", - "St. Martin (French part)", - "Turks and Caicos Islands", - "Virgin Islands (U.S.)", - "West Bank and Gaza", -) -non_sovereign_countries = list(NON_SOVEREIGN_COUNTRIES) - -NON_SOVEREIGN_IDS: frozenset[str] = frozenset( - { - "ABW", - "ASM", - "BMU", - "CHI", - "CUW", - "CYM", - "FRO", - "GIB", - "GRL", - "GUM", - "HKG", - "IMN", - "MAC", - "MAF", - "MNP", - "NCL", - "PRI", - "PSE", - "PYF", - "SXM", - "TCA", - "VGB", - "VIR", - } -) - -NON_SOVEREIGN_ISO2: frozenset[str] = frozenset( - { - "AW", - "AS", - "BM", - "JG", - "CW", - "KY", - "FO", - "GI", - "GL", - "GU", - "HK", - "IM", - "MO", - "MF", - "MP", - "NC", - "PR", - "PS", - "PF", - "SX", - "TC", - "VG", - "VI", - } -) - -AGGREGATE_IDS: frozenset[str] = frozenset( - { - "AFE", - "AFR", - "AFW", - "ARB", - "BEA", - "BEC", - "BHI", - "BLA", - "BMN", - "BSS", - "CAA", - "CEA", - "CEB", - "CEU", - "CLA", - "CME", - "CSA", - "CSS", - "DEA", - "DEC", - "DLA", - "DMN", - "DNS", - "DSA", - "DSF", - "DSS", - "EAP", - "EAR", - "EAS", - "ECA", - "ECS", - "EMU", - "EUU", - "FXS", - "HIC", - "HPC", - "IBB", - "IBD", - "IBT", - "IDA", - "IDB", - "IDX", - "INX", - "LAC", - "LCN", - "LDC", - "LIC", - "LMC", - "LMY", - "LTE", - "MDE", - "MEA", - "MIC", - "MNA", - "NAC", - "NAF", - "NRS", - "NXS", - "OED", - "OSS", - "PRE", - "PSS", - "PST", - "RRS", - "SAS", - "SSA", - "SSF", - "SST", - "SXZ", - "TEA", - "TEC", - "TLA", - "TMN", - "TSA", - "TSS", - "UMC", - "WLD", - "XZN", - } -) - -AGGREGATE_ISO2: frozenset[str] = frozenset( - { - "1A", - "1W", - "4E", - "7E", - "8S", - "B1", - "B2", - "B3", - "B4", - "B6", - "B7", - "B8", - "C4", - "C5", - "C6", - "C7", - "C8", - "C9", - "D1", - "D2", - "D3", - "D4", - "D5", - "D6", - "D7", - "EU", - "F1", - "F6", - "M1", - "M2", - "N6", - "OE", - "S1", - "S2", - "S3", - "S4", - "T2", - "T3", - "T4", - "T5", - "T6", - "T7", - "V1", - "V2", - "V3", - "V4", - "XC", - "XD", - "XE", - "XF", - "XG", - "XH", - "XI", - "XJ", - "XL", - "XM", - "XN", - "XO", - "XP", - "XQ", - "XR", - "XS", - "XT", - "XU", - "XY", - "XZ", - "Z4", - "Z7", - "ZB", - "ZF", - "ZG", - "ZH", - "ZI", - "ZJ", - "ZQ", - "ZT", - } -) - -AGGREGATE_NAMES: frozenset[str] = frozenset( - { - "Africa", - "Africa Eastern and Southern", - "Africa Western and Central", - "Arab World", - "Caribbean small states", - "Central Europe and the Baltics", - "Early-demographic dividend", - "East Asia & Pacific", - "East Asia & Pacific (IBRD-only countries)", - "East Asia & Pacific (IDA & IBRD countries)", - "East Asia & Pacific (IDA-eligible countries)", - "East Asia & Pacific (excluding high income)", - "East Asia and the Pacific (IFC classification)", - "Euro area", - "Europe & Central Asia", - "Europe & Central Asia (IBRD-only countries)", - "Europe & Central Asia (IDA & IBRD countries)", - "Europe & Central Asia (IDA-eligible countries)", - "Europe & Central Asia (excluding high income)", - "Europe and Central Asia (IFC classification)", - "European Union", - "Heavily indebted poor countries (HIPC)", - "High income", - "IBRD countries classified as high income", - "IBRD only", - "IBRD, including blend", - "IDA & IBRD total", - ( - "IDA countries classified as fragile situations, " - "excluding Sub-Saharan Africa" - ), - "IDA countries in Sub-Saharan Africa classified as fragile situations ", - "IDA countries in Sub-Saharan Africa not classified as fragile situations ", - ( - "IDA countries not classified as fragile situations, " - "excluding Sub-Saharan Africa" - ), - "IDA only", - "IDA total", - "Late-demographic dividend", - "Latin America & Caribbean ", - "Latin America & Caribbean (excluding high income)", - "Latin America & the Caribbean (IBRD-only countries)", - "Latin America & the Caribbean (IDA & IBRD countries)", - "Latin America & the Caribbean (IDA-eligible countries)", - "Latin America and the Caribbean (IFC classification)", - "Least developed countries: UN classification", - "Low & middle income", - "Low income", - "Lower middle income", - "Middle East (developing only)", - "Middle East and North Africa (IFC classification)", - "Middle East, North Africa, Afghanistan & Pakistan", - "Middle East, North Africa, Afghanistan & Pakistan (IBRD only)", - "Middle East, North Africa, Afghanistan & Pakistan (IDA & IBRD)", - "Middle East, North Africa, Afghanistan & Pakistan (IDA total)", - "Middle East, North Africa, Afghanistan & Pakistan (excluding high income)", - "Middle income", - "Non-resource rich Sub-Saharan Africa countries", - "North Africa", - "North America", - "Not classified", - "OECD members", - "Other small states", - "Pacific island small states", - "Post-demographic dividend", - "Pre-demographic dividend", - "Resource rich Sub-Saharan Africa countries", - "Small states", - "South Asia", - "South Asia (IDA & IBRD)", - "South Asia (IDA-eligible countries)", - "South Asia (IFC classification)", - "Sub-Saharan Africa ", - "Sub-Saharan Africa (IBRD-only countries)", - "Sub-Saharan Africa (IDA & IBRD countries)", - "Sub-Saharan Africa (IDA-eligible countries)", - "Sub-Saharan Africa (IFC classification)", - "Sub-Saharan Africa (excluding high income)", - "Sub-Saharan Africa excluding South Africa", - "Sub-Saharan Africa excluding South Africa and Nigeria", - "Upper middle income", - "World", - } -) - class SearchResult(list): """ @@ -483,91 +148,16 @@ def _filter_by_pattern( return (row for row in rows if pattern.search(row["name"])) -def _is_aggregate_entity(entity: dict[str, Any]) -> bool: - """Return True if an entity dict represents an aggregate region or group.""" - income_level = entity.get("incomeLevel") - if isinstance(income_level, dict) and income_level.get("value") == "Aggregates": - return True - region = entity.get("region") - if isinstance(region, dict) and ( - region.get("value") == "Aggregates" or region.get("id") == "NA" - ): - return True - lending_type = entity.get("lendingType") - if isinstance(lending_type, dict) and lending_type.get("value") == "Aggregates": - return True - if entity.get("id") in AGGREGATE_IDS: - return True - if entity.get("iso2Code") in AGGREGATE_ISO2: - return True - return entity.get("name") in AGGREGATE_NAMES - - -def _is_non_sovereign_entity(entity: dict[str, Any]) -> bool: - """Return True if an entity dict represents a non-sovereign country/territory.""" - if entity.get("name") in NON_SOVEREIGN_COUNTRIES: - return True - if entity.get("id") in NON_SOVEREIGN_IDS: - return True - return entity.get("iso2Code") in NON_SOVEREIGN_ISO2 - - -def _is_aggregate_row(row: dict[str, Any]) -> bool: - """Return True if an observation row represents an aggregate.""" - iso3 = row.get("countryiso3code") - if iso3 and iso3 in AGGREGATE_IDS: - return True - country_obj = row.get("country") - if isinstance(country_obj, dict): - if country_obj.get("id") in AGGREGATE_ISO2: - return True - if country_obj.get("value") in AGGREGATE_NAMES: - return True - return False - - -def _is_non_sovereign_row(row: dict[str, Any]) -> bool: - """Return True if an observation row represents a non-sovereign country.""" - iso3 = row.get("countryiso3code") - if iso3 and iso3 in NON_SOVEREIGN_IDS: - return True - country_obj = row.get("country") - if isinstance(country_obj, dict): - if country_obj.get("value") in NON_SOVEREIGN_COUNTRIES: - return True - if country_obj.get("id") in NON_SOVEREIGN_ISO2: - return True +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 -def _resolve_country_and_flags( - country: str | Sequence[str], - include_aggregates: bool, - include_non_sovereign: bool, - sovereign_only: bool, - exclude_aggregates: bool | None, -) -> tuple[str | Sequence[str], bool, bool]: - """Resolve and normalize country and filter flags.""" - if isinstance(country, str): - c_lower = country.strip().lower() - if c_lower in ("countries", "countries_only"): - country = "all" - include_aggregates = False - elif c_lower == "sovereign": - country = "all" - include_aggregates = False - include_non_sovereign = False - - if exclude_aggregates is not None: - include_aggregates = not exclude_aggregates - - if sovereign_only: - include_aggregates = False - include_non_sovereign = False - - return country, include_aggregates, include_non_sovereign - - @dataclasses.dataclass class Client: """ @@ -611,19 +201,14 @@ def get_data( source: int | str | Sequence[int | str] | None = None, parse_dates: bool = False, skip_cache: bool = False, - include_aggregates: bool = True, - include_non_sovereign: bool = True, - sovereign_only: bool = False, - exclude_aggregates: bool | None = None, ) -> fetcher.Result: """ Retrieve indicators for given countries and years Parameters: indicator: the desired indicator code - country: a country code, sequence of country codes, or "all" (default). - Also accepts special strings "countries" / "countries_only" (exclude - aggregates) and "sovereign" (sovereign countries only). + 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' @@ -634,26 +219,17 @@ def get_data( parse_dates: if True, convert date field to a datetime.datetime object. skip_cache: bypass the cache when downloading - include_aggregates: if False, exclude aggregate regions and groups - from the returned data. Defaults to True. - include_non_sovereign: if False, exclude non-sovereign - countries/territories from the returned data. Defaults to True. - sovereign_only: if True, exclude both aggregates and non-sovereign - countries. Defaults to False. - exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: A list of dictionaries of observations """ - country, include_aggregates, include_non_sovereign = ( - _resolve_country_and_flags( - country=country, - include_aggregates=include_aggregates, - include_non_sovereign=include_non_sovereign, - sovereign_only=sovereign_only, - exclude_aggregates=exclude_aggregates, - ) - ) + 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) @@ -668,19 +244,6 @@ def get_data( data = self.fetcher.fetch(url=url, params=params, skip_cache=skip_cache) if parse_dates: dates.parse_row_dates(data) - if (not include_aggregates or not include_non_sovereign) and isinstance( - data, list - ): - filtered_rows = [ - row - for row in data - if (include_aggregates or not _is_aggregate_row(row)) - and (include_non_sovereign or not _is_non_sovereign_row(row)) - ] - data = fetcher.Result( - filtered_rows, - last_updated=getattr(data, "last_updated", None), - ) return data def _id_only_query(self, url: str, id_: Any, skip_cache: bool) -> SearchResult: @@ -777,11 +340,8 @@ 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, - include_aggregates: bool = True, - include_non_sovereign: bool = True, - sovereign_only: bool = False, - exclude_aggregates: bool | None = None, ) -> SearchResult: """ Retrieve information on one or more country or regional aggregates. @@ -800,25 +360,14 @@ 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 - include_aggregates: if False, exclude aggregate regions and groups - from results. Defaults to True. - include_non_sovereign: if False, exclude non-sovereign - countries/territories from results. Defaults to True. - sovereign_only: if True, exclude both aggregates and non-sovereign - countries. Defaults to False. - exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: list of dictionaries describing countries """ - if exclude_aggregates is not None: - include_aggregates = not exclude_aggregates - if sovereign_only: - include_aggregates = False - include_non_sovereign = False - raw_results: Iterable[dict[str, Any]] if country_id: if incomelevel or lendingtype or query: @@ -837,13 +386,8 @@ def get_countries( ) raw_results = _filter_by_pattern(fetched, query) if query else fetched - if not include_aggregates or not include_non_sovereign: - raw_results = [ - row - for row in raw_results - if (include_aggregates or not _is_aggregate_entity(row)) - and (include_non_sovereign or not _is_non_sovereign_entity(row)) - ] + if exclude_aggregates: + raw_results = [r for r in raw_results if not _is_aggregate_country(r)] return SearchResult(raw_results) @@ -912,10 +456,6 @@ def get_series( name: str = "value", keep_levels: bool = False, skip_cache: bool = False, - include_aggregates: bool = True, - include_non_sovereign: bool = True, - sovereign_only: bool = False, - exclude_aggregates: bool | None = None, ) -> Series: """ Retrieve data for a single indicator as a pandas Series. @@ -924,9 +464,8 @@ def get_series( Parameters: indicator: the desired indicator code - country: a country code, sequence of country codes, or "all" (default). - Also accepts special strings "countries" / "countries_only" (exclude - aggregates) and "sovereign" (sovereign countries only). + 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' @@ -936,17 +475,11 @@ def get_series( to 2, World Development Indicators) parse_dates: if True, convert date field to a datetime.datetime object. + skip_cache: bypass the cache when downloading name: the desired name for the pandas Series 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 - include_aggregates: if False, exclude aggregate regions and groups - from the returned data. Defaults to True. - include_non_sovereign: if False, exclude non-sovereign - countries/territories from the returned data. Defaults to True. - sovereign_only: if True, exclude both aggregates and non-sovereign - countries. Defaults to False. - exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: Series with the requested data. The index of the series depends on @@ -958,16 +491,6 @@ def get_series( will be dropped. If `keep_levels` is `False` and both levels only have one value, the country level is dropped. """ - extra_kwargs: dict[str, Any] = {} - if not include_aggregates: - extra_kwargs["include_aggregates"] = include_aggregates - if not include_non_sovereign: - extra_kwargs["include_non_sovereign"] = include_non_sovereign - if sovereign_only: - extra_kwargs["sovereign_only"] = sovereign_only - if exclude_aggregates is not None: - extra_kwargs["exclude_aggregates"] = exclude_aggregates - raw_data = self.get_data( indicator=indicator, country=country, @@ -976,7 +499,6 @@ def get_series( source=source, parse_dates=parse_dates, skip_cache=skip_cache, - **extra_kwargs, ) df = pd.DataFrame( [[i["country"]["value"], i["date"], i["value"]] for i in raw_data], @@ -989,7 +511,7 @@ def get_series( df = df.set_index("country") else: df = df.set_index(["country", "date"]) - return Series(df[name], last_updated=getattr(raw_data, "last_updated", None)) + return Series(df[name], last_updated=raw_data.last_updated) @needs_pandas def get_dataframe( @@ -1005,22 +527,17 @@ def get_dataframe( parse_dates: bool = False, keep_levels: bool = False, skip_cache: bool = False, - include_aggregates: bool = True, - include_non_sovereign: bool = True, - sovereign_only: bool = False, - exclude_aggregates: bool | None = None, ) -> DataFrame: """ - Download a set of indicators and merge them into a pandas DataFrame. + Download a set of indicators and merge them into a pandas DataFrame. If pandas is not installed, a RuntimeError will be raised. Parameters: - indicators: A dictionary where the keys are desired indicators and the + 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). - Also accepts special strings "countries" / "countries_only" (exclude - aggregates) and "sovereign" (sovereign countries only). + 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' @@ -1033,13 +550,6 @@ def get_dataframe( 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 - include_aggregates: if False, exclude aggregate regions and groups - from the returned data. Defaults to True. - include_non_sovereign: if False, exclude non-sovereign - countries/territories from the returned data. Defaults to True. - sovereign_only: if True, exclude both aggregates and non-sovereign - countries. Defaults to False. - exclude_aggregates: if specified, sets include_aggregates to its negation. Returns: DataFrame with one column per indicator. The index of the DataFrame @@ -1052,16 +562,13 @@ def get_dataframe( levels only have one value, the country level is dropped. """ - extra_kwargs: dict[str, Any] = {} - if not include_aggregates: - extra_kwargs["include_aggregates"] = include_aggregates - if not include_non_sovereign: - extra_kwargs["include_non_sovereign"] = include_non_sovereign - if sovereign_only: - extra_kwargs["sovereign_only"] = sovereign_only - if exclude_aggregates is not None: - extra_kwargs["exclude_aggregates"] = exclude_aggregates - + 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( @@ -1073,18 +580,12 @@ def get_dataframe( parse_dates=parse_dates, keep_levels=True, skip_cache=skip_cache, - **extra_kwargs, ) for indicator, name in indicators.items() } ) - if ( - not keep_levels - and isinstance(df.index, pd.MultiIndex) - and len(df.index) > 0 - ): - if len(set(df.index.get_level_values(0))) == 1: - df.index = df.index.droplevel(0) - elif len(set(df.index.get_level_values(1))) == 1: - df.index = df.index.droplevel(1) + if not keep_levels and len(set(df.index.get_level_values(0))) == 1: + df.index = df.index.droplevel(0) + elif not keep_levels and len(set(df.index.get_level_values(1))) == 1: + df.index = df.index.droplevel(1) return df