correct absorption rester + add test

Author: esoteric-ephemera

mp_api/client/routes/materials/absorption.py

@@ -16,10 +16,11 @@ class AbsorptionRester(BaseRester):
     def search(
         self,
         material_ids: str | list[str] | None = None,
-        chemsys: str | list[str] | None = None,
-        elements: list[str] | None = None,
-        exclude_elements: list[str] | None = None,
-        formula: list[str] | None = None,
+        num_sites: int | tuple[int, int] | None = None,
+        num_elements: int | tuple[int, int] | None = None,
+        volume: float | tuple[float, float] | None = None,
+        density: float | tuple[float, float] | None = None,
+        band_gap: float | tuple[float, float] | None = None,
         num_chunks: int | None = None,
         chunk_size: int = 1000,
         all_fields: bool = True,
@@ -29,13 +30,6 @@ def search(
 
         Arguments:
             material_ids (str, List[str]): Search for optical absorption data associated with the specified Material IDs
-            chemsys (str, List[str]): A chemical system or list of chemical systems
-                (e.g., Li-Fe-O, Si-*, [Si-O, Li-Fe-P]).
-            elements (List[str]): A list of elements.
-            exclude_elements (List[str]): A list of elements to exclude.
-            formula (str, List[str]): A formula including anonymized formula
-                or wild cards (e.g., Fe2O3, ABO3, Si*). A list of chemical formulas can also be passed
-                (e.g., [Fe2O3, ABO3]).
             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.
@@ -46,23 +40,25 @@ def search(
         """
         query_params = defaultdict(dict)  # type: dict
 
-        if formula:
-            if isinstance(formula, str):
-                formula = [formula]
-
-            query_params.update({"formula": ",".join(formula)})
-
-        if chemsys:
-            if isinstance(chemsys, str):
-                chemsys = [chemsys]
-
-            query_params.update({"chemsys": ",".join(chemsys)})
-
-        if elements:
-            query_params.update({"elements": ",".join(elements)})
-
-        if exclude_elements:
-            query_params.update({"exclude_elements": ",".join(exclude_elements)})
+        aliased = {
+            "num_sites": "nsites",
+            "num_elements": "nelements",
+            "band_gap": "bandgap",
+        }
+        user_query = locals()
+        for k in ("num_sites","num_elements","volume","density","band_gap"):
+            if (value := user_query.get(k)) is not None:
+                if k in ("num_sites","num_elements") and isinstance(value, int):
+                    value = (value, value)
+                elif k in ("volume","density","band_gap") and isinstance(value,int | float):
+                    value = (value - 1e-2, value + 1e-2)
+
+                query_params.update(
+                    {
+                        f"{aliased.get(k,k)}_min": value[0],
+                        f"{aliased.get(k,k)}_max": value[1],
+                    }
+                )
 
         if material_ids:
             if isinstance(material_ids, str):
@@ -77,7 +73,6 @@ def search(
         }
 
         return super()._search(
-            formulae=formula,
             num_chunks=num_chunks,
             chunk_size=chunk_size,
             all_fields=all_fields,

mp_api/client/routes/materials/summary.py

@@ -48,8 +48,8 @@ def search(  # noqa: D417
         magnetic_ordering: Ordering | None = None,
         material_ids: str | list[str] | None = None,
         n: tuple[float, float] | None = None,
-        num_elements: tuple[int, int] | None = None,
-        num_sites: tuple[int, int] | None = None,
+        num_elements: int | tuple[int, int] | None = None,
+        num_sites: int | tuple[int, int] | None = None,
         num_magnetic_sites: tuple[int, int] | None = None,
         num_unique_magnetic_sites: tuple[int, int] | None = None,
         piezoelectric_modulus: tuple[float, float] | None = None,

tests/materials/test_absorption.py

@@ -0,0 +1,44 @@
+import os
+
+import pytest
+
+from emmet.core.phonon import PhononBS, PhononDOS
+
+from mp_api.client.core import MPRestError
+from mp_api.client.routes.materials.absorption import AbsorptionRester
+
+from core_function import client_search_testing
+
+@pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
+def test_absorption_search():
+
+    client_search_testing(
+        search_method=AbsorptionRester().search,
+        excluded_params=[
+            "num_chunks",
+            "chunk_size",
+            "all_fields",
+            "fields",
+        ],
+        alt_name_dict={
+            "material_ids": "material_id",
+            "num_sites": "nsites",
+            "num_elements": "nelements",
+            "band_gap": "bandgap",
+        },
+        custom_field_tests={
+            "material_ids": ["mp-149", "mp-239"],
+            "material_ids": "mp-149",
+            "num_sites": (6,7),
+            "num_sites": 7,
+            "num_elements": 5,
+            "num_elements": (4,5),
+            "volume": (115,116),
+            "volume": 115.5,
+            "density": (2.9,3),
+            "density": 2.933,
+            "band_gap": (1,1.05),
+            "band_gap": 1.,
+        },
+        sub_doc_fields=[],
+    )

tests/materials/test_doi.py

@@ -0,0 +1,28 @@
+import os
+
+import pytest
+
+from mp_api.client.routes.materials import DOIRester
+
+from core_function import client_search_testing
+
+@pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
+def test_doi_search():
+
+    client_search_testing(
+        search_method=DOIRester().search,
+        excluded_params=[
+            "num_chunks",
+            "chunk_size",
+            "all_fields",
+            "fields",
+        ],
+        alt_name_dict={
+            "material_ids": "material_id",
+        },
+        custom_field_tests={
+            "material_ids": ["mp-149", "mp-13"],
+            "material_ids": "mp-149",
+        },
+        sub_doc_fields=[],
+    )