From 4c8e928f167706183a4c5525e42970c0324890dd Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Wed, 21 Oct 2026 18:55:19 +0000 Subject: [PATCH] fix: keep columns that share a title instead of overwriting them get_data keyed renamed columns by title, so when two columns share a title (85592NED has "Transporttarief" under both Aardgas and Elektriciteit) the later one silently replaced the earlier: the gas tariff columns came back holding electricity values. columns= had the same collision and selected the electricity column. Shared titles are now prefixed with their topic group, falling back to the CBS key, and columns= raises on a shared title rather than guessing. Co-Authored-By: Claude Opus 5.5 --- docs/getting-started.md | 2 ++ src/cbspy/client.py | 29 +++++++++++++++- tests/test_client.py | 74 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 1 deletion(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 507239a..821ebcc 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -95,6 +95,8 @@ shape: (2, 28) Period codes like `2022JJ00` are automatically decoded to `2022`. Quarterly periods (`2023KW01`) become `2023 Q1`, and monthly periods (`2023MM03`) become `2023 March`. +Some tables use the same label for several columns. For example, 85592NED has "Transporttarief" under both Aardgas and Elektriciteit. cbspy then prefixes each label with its topic group: `Aardgas: Transporttarief` and `Elektriciteit: Transporttarief`. If the labels are still ambiguous, it appends the CBS key: `Totaal (Totaal_1)`. + ## Filtering by dimension By default, `get_data()` downloads the entire table. Use `periods` to filter by time, or `filters` to filter on any dimension: diff --git a/src/cbspy/client.py b/src/cbspy/client.py index 1ddca28..e5f3d3c 100644 --- a/src/cbspy/client.py +++ b/src/cbspy/client.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import Counter from typing import Any import httpx @@ -112,7 +113,7 @@ def get_data( Polars DataFrame with resolved column names and decoded periods. """ prop_rows = self._odata.get_json(table_id, "DataProperties") - column_map = {p["Key"]: p.get("Title", p["Key"]) for p in prop_rows} + column_map = self._column_names(prop_rows) period_keys = {p["Key"] for p in prop_rows if p.get("Type") == "TimeDimension"} filter_str = self._build_filter(periods, filters, period_keys) @@ -122,6 +123,11 @@ def get_data( if columns is not None: title_to_key = {v: k for k, v in column_map.items()} + for c in columns: + shared = [column_map[p["Key"]] for p in prop_rows if p.get("Key") and p.get("Title") == c] + if c not in title_to_key and c not in column_map and shared: + msg = f"{c!r} is the title of several columns in {table_id}; use one of {shared}" + raise ValueError(msg) select_keys = [title_to_key.get(c, c) for c in columns] params["$select"] = ",".join(select_keys) @@ -146,6 +152,27 @@ def get_data( return pl.DataFrame(renamed_rows) + @staticmethod + def _column_names(prop_rows: list[dict[str, Any]]) -> dict[str, str]: + """Map each column key to its title, qualifying titles that several columns share. + + A shared title is prefixed with its topic group ("Aardgas: Transporttarief"). If that is + still ambiguous, the CBS key is appended ("Totaal (Totaal_1)"). + """ + groups = {p.get("ID"): p.get("Title") or "" for p in prop_rows if p.get("Type") == "TopicGroup"} + cols = [p for p in prop_rows if p.get("Key")] + titles = Counter(p.get("Title") or p["Key"] for p in cols) + + names: dict[str, str] = {} + for p in cols: + name = p.get("Title") or p["Key"] + if titles[name] > 1 and p.get("ParentID") in groups: + name = f"{groups[p['ParentID']]}: {name}" + names[p["Key"]] = name + + counts = Counter(names.values()) + return {k: f"{v} ({k})" if counts[v] > 1 else v for k, v in names.items()} + @staticmethod def _build_filter( periods: list[str] | None, diff --git a/tests/test_client.py b/tests/test_client.py index b4ca533..b3afc8c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -571,6 +571,80 @@ def test_typed_default_is_true(self): assert route.called +GROUPED_PROPS = [ + {"ID": 0, "Key": "Perioden", "Title": "Perioden", "Type": "TimeDimension"}, + {"ID": 1, "Key": "", "Title": "Aardgas", "Type": "TopicGroup"}, + {"ID": 2, "ParentID": 1, "Key": "Transporttarief_1", "Title": "Transporttarief", "Type": "Topic"}, + {"ID": 3, "Key": "", "Title": "Elektriciteit", "Type": "TopicGroup"}, + {"ID": 4, "ParentID": 3, "Key": "Transporttarief_7", "Title": "Transporttarief", "Type": "Topic"}, + {"ID": 5, "ParentID": 3, "Key": "Vermindering_15", "Title": "Vermindering energiebelasting", "Type": "Topic"}, +] + + +class TestGetDataDuplicateTitles: + """Columns that share a title, e.g. 85592NED has "Transporttarief" under both Aardgas and Elektriciteit.""" + + def _mock_props(self, props: list[dict]) -> None: + respx.get(f"{BASE}/ODataFeed/OData/85592NED/DataProperties").mock( + return_value=httpx.Response(200, json={"value": props}) + ) + + def _mock_rows(self, rows: list[dict]) -> respx.Route: + return respx.get(f"{BASE}/ODataFeed/OData/85592NED/TypedDataSet").mock( + return_value=httpx.Response(200, json={"value": rows}) + ) + + @respx.mock + def test_shared_titles_are_qualified_by_group(self): + self._mock_props(GROUPED_PROPS) + self._mock_rows([ + { + "ID": 0, + "Perioden": "2026MM09", + "Transporttarief_1": 221.79, + "Transporttarief_7": 394.11, + "Vermindering_15": -519.8, + } + ]) + with Client() as client: + df = client.get_data("85592NED") + assert df.columns == [ + "Perioden", + "Aardgas: Transporttarief", + "Elektriciteit: Transporttarief", + "Vermindering energiebelasting", + ] + assert df["Aardgas: Transporttarief"][0] == 221.79 + assert df["Elektriciteit: Transporttarief"][0] == 394.11 + + @respx.mock + def test_shared_titles_without_group_fall_back_to_key(self): + self._mock_props([ + {"Key": "Perioden", "Title": "Perioden", "Type": "TimeDimension"}, + {"Key": "Totaal_1", "Title": "Totaal", "Type": "Topic"}, + {"Key": "Totaal_2", "Title": "Totaal", "Type": "Topic"}, + ]) + self._mock_rows([{"ID": 0, "Perioden": "2026MM09", "Totaal_1": 1, "Totaal_2": 2}]) + with Client() as client: + df = client.get_data("85592NED") + assert df["Totaal (Totaal_1)"][0] == 1 + assert df["Totaal (Totaal_2)"][0] == 2 + + @respx.mock + def test_columns_accepts_qualified_title(self): + self._mock_props(GROUPED_PROPS) + route = self._mock_rows([{"ID": 0, "Perioden": "2026MM09", "Transporttarief_1": 221.79}]) + with Client() as client: + client.get_data("85592NED", columns=["Perioden", "Aardgas: Transporttarief"]) + assert dict(route.calls[0].request.url.params)["$select"] == "Perioden,Transporttarief_1" + + @respx.mock + def test_columns_rejects_shared_title(self): + self._mock_props(GROUPED_PROPS) + with Client() as client, pytest.raises(ValueError, match="Aardgas: Transporttarief"): + client.get_data("85592NED", columns=["Transporttarief"]) + + class TestClientLifecycle: def test_close_closes_owned_http_client(self): client = Client()