fix eos, task traj, add __dir__ to lazy import + fix attr access

Author: esoteric-ephemera

mp_api/_test_utils.py

@@ -78,6 +78,7 @@ def client_search_testing(
                 doc = docs[0].model_dump()
             else:
                 raise ValueError("No documents returned")
+            print(doc)
 
             for sub_field in sub_doc_fields:
                 if sub_field in doc:

mp_api/client/core/client.py

@@ -23,6 +23,7 @@
 from itertools import chain, islice
 from json import JSONDecodeError
 from math import ceil
+from pathlib import Path
 from typing import TYPE_CHECKING, ForwardRef, Optional, get_args
 from urllib.parse import urljoin
 
@@ -184,7 +185,7 @@ def __init__(
 
         self.use_document_model = use_document_model
         self.mute_progress_bars = mute_progress_bars
-        self.local_dataset_cache = local_dataset_cache
+        self.local_dataset_cache = Path(local_dataset_cache)
         self.force_renew = force_renew
         self._query_builder = query_builder
 
@@ -1436,12 +1437,7 @@ def _convert_to_model(
             )
 
             return [
-                data_model(
-                    **{
-                        field: raw_doc[field]
-                        for field in set_fields.intersection(raw_doc)
-                    }
-                )
+                data_model(**raw_doc)
                 for raw_doc in (data if is_list else chain([first_doc], data))
             ]
 
@@ -1464,7 +1460,14 @@ def _generate_returned_model(
             set of str: set_fields, fields_not_requested)
         """
         model_fields = self.document_model.model_fields
-        set_fields = set(doc).intersection(model_fields)
+        aliases = {
+            anno.alias: field for field, anno in model_fields.items() if anno.alias
+        }
+        set_fields = (
+            set(doc)
+            .intersection(model_fields)
+            .union({aliases[k] for k in set(doc).intersection(aliases)})
+        )
         unset_fields = set(model_fields).difference(set_fields)
         user_requested_fields: list[str] = requested_fields or []
         fields_not_requested = unset_fields.difference(user_requested_fields)

mp_api/client/core/utils.py

@@ -116,8 +116,8 @@ def validate_ids(id_list: list[str]) -> list[str]:
             " data for all IDs and filter locally."
         )
 
-    [validate_identifier(idx, serialize=False) for idx in id_list]
-    return [getattr(idx, "string", str(idx)) for idx in id_list]
+    validated = [validate_identifier(idx, serialize=False) for idx in id_list]
+    return [getattr(idx, "string", str(idx)) for idx in validated]
 
 
 def validate_endpoint(endpoint: str | None, suffix: str | None = None) -> str:
@@ -241,6 +241,14 @@ def __getattr__(self, v: str) -> Any:
         if hasattr(self._imported, v):
             return getattr(self._imported, v)
 
+        raise AttributeError(
+            f"{self._module_name}{'.' + self._class_name if self._class_name else ''} "
+            f"has no attribute {v}"
+        )
+
+    def __dir__(self) -> list[str]:
+        return self._obj.__dir__()
+
 
 class MPDataset:
     """Convenience wrapper for pyarrow datasets stored on disk."""

mp_api/client/routes/materials/electronic_structure.py

@@ -326,7 +326,9 @@ def get_bandstructure_from_material_id(
                 material_ids=material_id, fields=["bandstructure"]
             )
             if not bs_doc:
-                raise MPRestError("No electronic structure data found.")
+                raise MPRestError(
+                    f"No electronic structure data found for material ID {material_id}."
+                )
 
             if (_bs_data := bs_doc[0]["bandstructure"]) is None:
                 raise MPRestError(
@@ -349,7 +351,9 @@ def get_bandstructure_from_material_id(
                     material_ids=material_id, fields=["dos"]
                 )
             ):
-                raise MPRestError("No electronic structure data found.")
+                raise MPRestError(
+                    f"No electronic structure data found for material ID {material_id}."
+                )
 
             if (_bs_data := bs_doc[0]["dos"]) is None:
                 raise MPRestError(
@@ -538,7 +542,9 @@ def get_dos_from_material_id(self, material_id: str) -> Dos:
         if not (
             dos_doc := self.es_rester.search(material_ids=material_id, fields=["dos"])
         ):
-            return None
+            raise MPRestError(
+                f"No electronic structure data found for material ID {material_id}."
+            )
 
         if not (dos_data := dos_doc[0].get("dos")):
             raise MPRestError(f"No density of states data found for {material_id}")

mp_api/client/routes/materials/eos.py

@@ -1,50 +1,66 @@
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 
 from emmet.core.eos import EOSDoc
 
-from mp_api.client.core import BaseRester
+from mp_api.client.core import BaseRester, MPRestError, MPRestWarning
 from mp_api.client.core.utils import validate_ids
 
 
 class EOSRester(BaseRester):
     suffix = "materials/eos"
     document_model = EOSDoc  # type: ignore
-    primary_key = "material_id"
+    primary_key = "task_id"
 
     def search(
         self,
-        material_ids: str | list[str] | None = None,
+        task_ids: str | list[str] | None = None,
         energies: tuple[float, float] | None = None,
         volumes: tuple[float, float] | None = None,
         num_chunks: int | None = None,
         chunk_size: int = 1000,
         all_fields: bool = True,
         fields: list[str] | None = None,
+        **kwargs,
     ) -> list[EOSDoc] | list[dict]:
         """Query equations of state docs using a variety of search criteria.
 
         Arguments:
-            material_ids (str, List[str]): Search for equation of states associated with the specified Material IDs
+            task_ids (str, List[str]): Search for equation of states associated with the specified task IDs
             energies (Tuple[float,float]): Minimum and maximum energy in eV/atom to consider for EOS plot range.
             volumes (Tuple[float,float]): Minimum and maximum volume in A³/atom to consider for EOS plot range.
             num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
             chunk_size (int): Number of data entries per chunk.
             all_fields (bool): Whether to return all fields in the document. Defaults to True.
             fields (List[str]): List of fields in EOSDoc to return data for.
                 Default is material_id only if all_fields is False.
+            **kwargs : used for handling deprecated kwargs
 
         Returns:
             ([EOSDoc], [dict]) List of equations of state docs or dictionaries.
         """
         query_params: dict = defaultdict(dict)
 
-        if material_ids:
-            if isinstance(material_ids, str):
-                material_ids = [material_ids]
+        if "material_ids" in kwargs:
+            if task_ids:
+                raise MPRestError(
+                    "You have specified both `task_ids` and the deprecated `material_ids` tag. "
+                    "Please specify only `task_ids`."
+                )
+            task_ids = kwargs.pop("material_ids")
+            warnings.warn(
+                "`material_id` has been replaced by `task_id` in the EOS endpoint. "
+                "Please migrate to using the newer field name.",
+                stacklevel=2,
+                category=MPRestWarning,
+            )
 
-            query_params.update({"material_ids": ",".join(validate_ids(material_ids))})
+        if task_ids:
+            query_params["material_ids"] = ",".join(
+                validate_ids([task_ids] if isinstance(task_ids, str) else task_ids)
+            )
 
         if volumes:
             query_params.update({"volumes_min": volumes[0], "volumes_max": volumes[1]})

mp_api/client/routes/materials/tasks.py

@@ -39,10 +39,8 @@ def get_trajectory(
         """
         as_alpha = str(AlphaID(task_id, padlen=8)).split("-")[-1]
         predicate = (
-            f"WHERE run_type='{str(run_type)}' AND identifier='{as_alpha}'"
-            if run_type
-            else f"WHERE identifier='{as_alpha}'"
-        )
+            f"WHERE run_type='{str(run_type)}' AND " if run_type else ""
+        ) + f"WHERE identifier='{as_alpha}'"
 
         traj_lbl, traj_tbl = self._get_delta_table(
             "materialsproject-parsed",
@@ -53,7 +51,6 @@ def get_trajectory(
         query = f"""
             SELECT *
             FROM   {traj_lbl}
-            WHERE  identifier='{as_alpha}'
             {predicate};
         """
 

tests/client/materials/test_electronic_structure.py

@@ -104,7 +104,7 @@ def test_bs_client():
         with pytest.raises(MPRestError, match="No electronic structure data found."):
             _ = bs_rester.get_bandstructure_from_material_id("mp-0")
 
-        with pytest.raises(MPRestError, match="No object found"):
+        with pytest.raises(MPRestError, match="No bandstructure data found"):
             _ = bs_rester.get_bandstructure_from_task_id("mp-0")
 
 

tests/client/materials/test_eos.py

@@ -4,6 +4,7 @@
 
 from mp_api._test_utils import client_search_testing, requires_api_key
 
+from mp_api.client.core.exceptions import MPRestError, MPRestWarning
 from mp_api.client.routes.materials.eos import EOSRester
 
 
@@ -26,9 +27,9 @@ def rester():
 
 sub_doc_fields: list = []
 
-alt_name_dict: dict = {"material_ids": "material_id"}
+alt_name_dict: dict = {"task_ids": "task_id"}
 
-custom_field_tests: dict = {"material_ids": ["mp-149"]}
+custom_field_tests: dict = {"task_ids": ["mp-149"]}
 
 
 @requires_api_key
@@ -42,3 +43,15 @@ def test_client(rester):
         custom_field_tests=custom_field_tests,
         sub_doc_fields=sub_doc_fields,
     )
+
+
+@requires_api_key
+def test_warnings_errors(rester):
+
+    with pytest.warns(
+        MPRestWarning, match="`material_id` has been replaced by `task_id`"
+    ):
+        rester.search(material_ids=["mp-149"], num_chunks=1, chunk_size=1)
+
+    with pytest.raises(MPRestError, match="You have specified both"):
+        rester.search(material_ids=["mp-149"], task_ids=["mp-1"])