Skip to content

Commit 1426b2e

Browse files
broken phase diagram serialization
1 parent 3c2a7b4 commit 1426b2e

50 files changed

Lines changed: 837 additions & 794 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/testing.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ jobs:
5858
run: python -m mypy mp_api/
5959

6060
- name: Test with pytest
61+
if: always()
6162
env:
6263
MP_API_KEY: ${{ secrets[env.API_KEY_NAME] }}
6364
# MP_API_ENDPOINT: https://api-preview.materialsproject.org/

mp_api/_test_utils.py

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
"""Define testing utils that need to imported."""
2+
3+
# pragma: exclude file
4+
5+
from __future__ import annotations
6+
7+
try:
8+
import pytest
9+
except ImportError as exc:
10+
raise ImportError(
11+
"You must `pip install 'mp-api[test]' to use these testing utilities."
12+
) from exc
13+
14+
import os
15+
from typing import TYPE_CHECKING
16+
17+
if TYPE_CHECKING:
18+
from collections.abc import Callable, Sequence
19+
from typing import Any
20+
21+
requires_api_key = pytest.mark.skipif(
22+
os.getenv("MP_API_KEY") is None,
23+
reason="No API key found.",
24+
)
25+
26+
NUM_DOCS = 5
27+
28+
29+
def client_search_testing(
30+
search_method: Callable,
31+
excluded_params: list[str],
32+
alt_name_dict: dict[str, str],
33+
custom_field_tests: dict[str, Any],
34+
sub_doc_fields: list[str],
35+
int_bounds: tuple[int, int] = (-100, 100),
36+
float_bounds: tuple[float, float] = (-100.12, 100.12),
37+
):
38+
"""Function to test a client using its search method.
39+
Each parameter is used to query for data, which is then checked.
40+
41+
Args:
42+
search_method (Callable): Client search method
43+
excluded_params (list[str]): List of parameters to exclude from testing
44+
alt_name_dict (dict[str, str]): Alternative names for parameters used in the projection and subsequent data checking
45+
custom_field_tests (dict[str, Any]): Custom queries for specific fields.
46+
sub_doc_fields (list[str]): Prefixes for fields to check in resulting data. Useful when data to be tested is nested.
47+
int_bounds (tuple[int,int]) : integer bounds to use in testing int-type query arguments
48+
float_bounds (tuple[float,float]) : float bounds to use in testing float-type query arguments
49+
"""
50+
if search_method is None:
51+
return
52+
# Get list of parameters
53+
param_tuples = list(search_method.__annotations__.items())
54+
55+
# Query API for each numeric and boolean parameter and check if returned
56+
for entry in param_tuples:
57+
param = entry[0]
58+
59+
if param not in excluded_params + ["return"]:
60+
param_type = entry[1]
61+
q: dict[str, Any] = {"chunk_size": 1, "num_chunks": 1}
62+
63+
if "tuple[int, int]" in param_type:
64+
q[param] = int_bounds
65+
elif "tuple[float, float]" in param_type:
66+
q[param] = float_bounds
67+
elif "bool" in param_type:
68+
q[param] = False
69+
elif param in custom_field_tests:
70+
q[param] = custom_field_tests[param]
71+
else:
72+
raise ValueError(
73+
f"Parameter '{param}' with type '{param_type}' was not "
74+
"properly identified in the generic search method test."
75+
)
76+
77+
if len(docs := search_method(**q)) > 0:
78+
doc = docs[0].model_dump()
79+
else:
80+
raise ValueError("No documents returned")
81+
82+
for sub_field in sub_doc_fields:
83+
if sub_field in doc:
84+
doc = doc[sub_field]
85+
86+
assert doc[alt_name_dict.get(param, param)] is not None
87+
88+
89+
def client_pagination(search_method: Callable, id_name: str):
90+
page_1 = search_method(_page=1, chunk_size=NUM_DOCS, fields=[id_name])
91+
page_2 = search_method(_page=2, chunk_size=NUM_DOCS, fields=[id_name])
92+
assert all(len(results) == NUM_DOCS for results in (page_1, page_2))
93+
assert {str(getattr(doc, id_name)) for doc in page_1}.intersection(
94+
{str(getattr(doc, id_name)) for doc in page_2}
95+
) == set()
96+
97+
98+
def client_sort(search_method: Callable, sort_fields: str | Sequence[str]):
99+
for sort_field in [sort_fields] if isinstance(sort_fields, str) else sort_fields:
100+
asc = search_method(
101+
_page=1, _sort_fields=sort_field, chunk_size=NUM_DOCS, fields=[sort_field]
102+
)
103+
desc = search_method(
104+
_page=1,
105+
_sort_fields=f"-{sort_field}",
106+
chunk_size=NUM_DOCS,
107+
fields=[sort_field],
108+
)
109+
110+
idxs = list(range(NUM_DOCS))
111+
assert sorted(idxs, key=lambda idx: getattr(asc[idx], sort_field)) == idxs
112+
113+
assert (
114+
sorted(
115+
idxs,
116+
key=lambda idx: getattr(desc[idx], sort_field),
117+
reverse=True,
118+
)
119+
== idxs
120+
)

mp_api/client/_server_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ def get_request_headers() -> dict[str, Any]:
3838
Empty dict if flask is not installed, or not in a request context.
3939
Request headers otherwise.
4040
"""
41-
return request.headers if has_request_context() else {}
41+
return dict(request.headers) if has_request_context() else {}
4242

4343

4444
def is_dev_env(

mp_api/client/core/_oxygen_evolution.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ def get_oxygen_evolution_from_phase_diagram(
167167
"evolution": "evolution",
168168
}.items()
169169
}
170-
for profile in phase_diagram.get_element_profile("O", composition)
170+
for profile in phase_diagram.get_element_profile("O", composition) # type: ignore[arg-type]
171171
]
172172
for composition in set(compositions)
173173
}

mp_api/client/core/client.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -952,6 +952,10 @@ def _submit_requests(
952952
# No splitting needed - get first page
953953
total_data = {"data": []}
954954
initial_criteria = copy(criteria)
955+
if isinstance(
956+
initial_criteria.get("_page"), int
957+
) and not initial_criteria.get("_per_page"):
958+
initial_criteria["_per_page"] = initial_criteria.get("_limit")
955959
data, total_num_docs = self._submit_request_and_process(
956960
url=url,
957961
verify=True,
@@ -1277,6 +1281,9 @@ def _search(
12771281
# This method should be customized for each end point to give more user friendly,
12781282
# documented kwargs.
12791283

1284+
# If user specifies page, ensure only one chunk is returned
1285+
if isinstance(kwargs.get("_page"), int) and num_chunks is None:
1286+
num_chunks = 1
12801287
return self._get_all_documents(
12811288
kwargs,
12821289
all_fields=all_fields,

mp_api/client/routes/materials/electrodes.py

Lines changed: 33 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ class BaseElectrodeRester(BaseRester):
1515
primary_key = "battery_id"
1616
_exclude_search_fields: list[str] | None = None
1717

18-
def search( # pragma: ignore
18+
def search(
1919
self,
2020
battery_ids: str | list[str] | None = None,
2121
average_voltage: tuple[float, float] | None = None,
@@ -39,6 +39,8 @@ def search( # pragma: ignore
3939
chunk_size: int = 1000,
4040
all_fields: bool = True,
4141
fields: list[str] | None = None,
42+
_page: int | None = None,
43+
_sort_fields: str | None = None,
4244
) -> list[InsertionElectrodeDoc | ConversionElectrodeDoc] | list[dict]:
4345
"""Query using a variety of search criteria.
4446
@@ -77,63 +79,45 @@ def search( # pragma: ignore
7779
all_fields (bool): Whether to return all fields in the document. Defaults to True.
7880
fields (List[str]): List of fields in InsertionElectrodeDoc or ConversionElectrodeDoc to return data for.
7981
Default is battery_id and last_updated if all_fields is False.
82+
_page (int or None) : Page of the results to skip to.
83+
_sort_fields (str or None) : Field to sort on. Including a leading "-" sign will reverse sort order.
8084
8185
Returns:
8286
([InsertionElectrodeDoc or ConversionElectrodeDoc], [dict]) List of insertion/conversion electrode documents or dictionaries.
8387
"""
8488
query_params: dict = defaultdict(dict)
8589

86-
if battery_ids:
87-
if isinstance(battery_ids, str):
88-
battery_ids = [battery_ids]
89-
90-
query_params.update({"battery_ids": ",".join(validate_ids(battery_ids))})
91-
92-
if working_ion:
93-
if isinstance(working_ion, (str, Element)):
94-
working_ion = [working_ion] # type: ignore
95-
96-
query_params.update(
97-
{"working_ion": ",".join([str(ele) for ele in working_ion])} # type: ignore
98-
)
99-
100-
if formula:
101-
if isinstance(formula, str):
102-
formula = [formula]
103-
104-
query_params.update({"formula": ",".join(formula)})
105-
106-
if elements:
107-
query_params.update({"elements": ",".join(elements)})
108-
109-
if num_elements:
110-
if isinstance(num_elements, int):
111-
num_elements = (num_elements, num_elements)
112-
query_params.update(
113-
{"nelements_min": num_elements[0], "nelements_max": num_elements[1]}
114-
)
115-
116-
if exclude_elements:
117-
query_params.update({"exclude_elements": ",".join(exclude_elements)})
118-
11990
for param, value in locals().items():
120-
if (
121-
param
122-
not in [
123-
"__class__",
124-
"self",
125-
"working_ion",
126-
"query_params",
127-
"num_elements",
128-
]
129-
and value
130-
):
131-
if isinstance(value, tuple):
91+
if param not in {"__class__", "self", "query_params"} and value is not None:
92+
if param == "num_elements": # this must come first
93+
if isinstance(num_elements, int):
94+
num_elements = (num_elements, num_elements)
95+
query_params.update(
96+
{
97+
"nelements_min": num_elements[0], # type: ignore[index]
98+
"nelements_max": num_elements[1], # type: ignore[index]
99+
}
100+
)
101+
102+
elif isinstance(value, tuple):
132103
query_params.update(
133104
{f"{param}_min": value[0], f"{param}_max": value[1]}
134105
)
106+
elif param == "battery_ids":
107+
query_params[param] = ",".join(validate_ids(value))
108+
elif param == "working_ion":
109+
query_params["working_ion"] = ",".join(
110+
str(ele)
111+
for ele in (
112+
[value] if isinstance(value, str | Element) else value
113+
)
114+
)
115+
elif param in ("formula", "elements", "exclude_elements"):
116+
query_params[param] = ",".join(
117+
[value] if isinstance(value, str) else value
118+
)
135119
else:
136-
query_params.update({param: value})
120+
query_params[param] = value
137121

138122
excluded_fields = self._exclude_search_fields or []
139123
ignored_fields = {
@@ -177,4 +161,6 @@ class ConversionElectrodeRester(BaseElectrodeRester):
177161
"stability_charge",
178162
"stability_discharge",
179163
"exclude_elements",
164+
"_page",
165+
"_sort_fields",
180166
]

0 commit comments

Comments
 (0)