Skip to content

Commit a2017fa

Browse files
authored
fetch_dcids_by_wikidata_id improvement (#227)
- Renames `.fetch_dcid_by_wikidata_id` to `fetch_dcids_by_wikidata_id` to clarify lists are accepted. - Adds a test for lists of `wikidata_ids`
1 parent 4ce786d commit a2017fa

File tree

2 files changed

+33
-12
lines changed

2 files changed

+33
-12
lines changed

datacommons_client/endpoints/resolve.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -81,25 +81,25 @@ def fetch_dcids_by_name(self,
8181

8282
return self.fetch(node_ids=names, expression=expression)
8383

84-
def fetch_dcid_by_wikidata_id(
84+
def fetch_dcids_by_wikidata_id(
8585
self,
86-
wikidata_id: str | list[str],
86+
wikidata_ids: str | list[str],
8787
entity_type: Optional[str] = None) -> ResolveResponse:
8888
"""
89-
Fetches DCIDs for entities by their Wikidata IDs.
89+
Fetches DCIDs for entities by their Wikidata IDs.
9090
91-
Args:
92-
wikidata_id (str | list[str]): One or more Wikidata IDs to resolve.
93-
entity_type (Optional[str]): Optional type of the entities.
91+
Args:
92+
wikidata_ids (str | list[str]): One or more Wikidata IDs to resolve.
93+
entity_type (Optional[str]): Optional type of the entities.
9494
95-
Returns:
96-
ResolveResponse: The response object containing the resolved DCIDs.
97-
"""
95+
Returns:
96+
ResolveResponse: The response object containing the resolved DCIDs.
97+
"""
9898
expression = _resolve_correspondence_expression(from_type="wikidataId",
9999
to_type="dcid",
100100
entity_type=entity_type)
101101

102-
return self.fetch(node_ids=wikidata_id, expression=expression)
102+
return self.fetch(node_ids=wikidata_ids, expression=expression)
103103

104104
def fetch_dcid_by_coordinates(
105105
self,

datacommons_client/tests/endpoints/test_resolve_endpoint.py

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,8 @@ def test_fetch_dcid_by_wikidata_id():
5252
api_mock = MagicMock(spec=API)
5353
endpoint = ResolveEndpoint(api=api_mock)
5454

55-
response = endpoint.fetch_dcid_by_wikidata_id(wikidata_id="Q12345",
56-
entity_type="Country")
55+
response = endpoint.fetch_dcids_by_wikidata_id(wikidata_ids="Q12345",
56+
entity_type="Country")
5757

5858
# Check the response
5959
assert isinstance(response, ResolveResponse)
@@ -68,6 +68,27 @@ def test_fetch_dcid_by_wikidata_id():
6868
next_token=None)
6969

7070

71+
def test_fetch_dcids_list_by_wikidata_id():
72+
"""Tests the fetch_dcid_by_wikidata_id method."""
73+
api_mock = MagicMock(spec=API)
74+
endpoint = ResolveEndpoint(api=api_mock)
75+
76+
response = endpoint.fetch_dcids_by_wikidata_id(
77+
wikidata_ids=["Q12345", "Q695660"])
78+
79+
# Check the response
80+
assert isinstance(response, ResolveResponse)
81+
82+
# Check the post request
83+
api_mock.post.assert_called_once_with(payload={
84+
"nodes": ["Q12345", "Q695660"],
85+
"property": "<-wikidataId->dcid",
86+
},
87+
endpoint="resolve",
88+
all_pages=True,
89+
next_token=None)
90+
91+
7192
def test_fetch_dcid_by_coordinates():
7293
"""Tests the fetch_dcid_by_coordinates method."""
7394
api_mock = MagicMock(spec=API)

0 commit comments

Comments
 (0)