update phonon rester to use new fields, maintain backwards compat with warnings

Author: tsmathis

mp_api/client/routes/materials/phonon.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 from typing import TYPE_CHECKING
 
 from emmet.core.phonon import PhononBS, PhononBSDOSDoc, PhononDOS
 
-from mp_api.client.core import BaseRester, MPRestError
+from mp_api.client.core import BaseRester, MPRestError, MPRestWarning
 from mp_api.client.core.utils import validate_ids
 
 if TYPE_CHECKING:
@@ -17,40 +18,58 @@
 class PhononRester(BaseRester):
     suffix = "materials/phonon"
     document_model = PhononBSDOSDoc  # type: ignore
-    primary_key = "material_id"
+    primary_key = "identifier"
     delta_backed = False
 
     def search(
         self,
-        material_ids: str | list[str] | None = None,
+        phonon_ids: str | list[str] | None = None,
         phonon_method: str | None = None,
         num_chunks: int | None = None,
         chunk_size: int = 1000,
         all_fields: bool = True,
         fields: list[str] | None = None,
+        **kwargs,
     ) -> list[PhononBSDOSDoc] | list[dict]:
         """Query phonon docs using a variety of search criteria.
 
         Arguments:
-            material_ids (str, List[str]): A single Material ID string or list of strings
+            phonon_ids (str, List[str]): A single Phonon Task ID string or list of strings
                 (e.g., mp-149, [mp-149, mp-13]).
             phonon_method (str): phonon method to search (dfpt, phonopy, pheasy)
             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 PhononBSDOSDoc to return data for.
                 Default is material_id, last_updated, and formula_pretty if all_fields is False.
+            **kwargs : used for handling deprecated kwargs
 
         Returns:
             ([PhononBSDOSDoc], [dict]) List of phonon documents or dictionaries.
         """
         query_params: dict = defaultdict(dict)
 
-        if material_ids:
-            if isinstance(material_ids, str):
-                material_ids = [material_ids]
-
-            query_params["material_ids"] = ",".join(validate_ids(material_ids))
+        if "material_ids" in kwargs:
+            if phonon_ids:
+                raise MPRestError(
+                    "You have specified both `phonon_ids` and the deprecated `material_ids` tag. "
+                    "Please specify only `phonon_ids`."
+                )
+            phonon_ids = kwargs.pop("material_ids")
+            warnings.warn(
+                "`material_id` has been replaced by `identifier` in the phonon endpoint. "
+                "Please migrate to using the newer field name and the more generic `phonon_ids` kwarg "
+                "for searching.",
+                stacklevel=2,
+                category=MPRestWarning,
+            )
+
+        if phonon_ids:
+            query_params["phonon_ids"] = ",".join(
+                validate_ids(
+                    [phonon_ids] if isinstance(phonon_ids, str) else phonon_ids
+                )
+            )
 
         if phonon_method and phonon_method in {"dfpt", "phonopy", "pheasy"}:
             query_params["phonon_method"] = phonon_method
@@ -69,21 +88,21 @@ def search(
             **query_params,
         )
 
-    def get_bandstructure_from_material_id(
-        self, material_id: str, phonon_method: str
+    def get_bandstructure_from_phonon_id(
+        self, phonon_id: str, phonon_method: str
     ) -> PhononBS | dict[str, Any]:
-        """Get the phonon band structure pymatgen object associated with a given material ID and phonon method.
+        """Get the phonon band structure pymatgen object associated with a given phonon ID and phonon method.
 
         Arguments:
-            material_id (str): Material ID for the phonon band structure calculation
+            phonon_id (str): Phonon ID for the phonon band structure calculation
             phonon_method (str): phonon method, i.e. pheasy or dfpt
 
         Returns:
             bandstructure (PhononBS): PhononBS object
         """
         result = self._query_open_data(
             bucket="materialsproject-parsed",
-            key=f"ph-bandstructures/{phonon_method}/{material_id}.json.gz",
+            key=f"ph-bandstructures/{phonon_method}/{phonon_id}.json.gz",
         )[0][0]
 
         return (
@@ -92,21 +111,33 @@ def get_bandstructure_from_material_id(
             else result  # type: ignore[return-value]
         )
 
-    def get_dos_from_material_id(
+    def get_bandstructure_from_material_id(
         self, material_id: str, phonon_method: str
+    ) -> PhononBS | dict[str, Any]:
+        """Deprecated: use `get_bandstructure_from_phonon_id` instead."""
+        warnings.warn(
+            "`material_id` has been replaced by `identifier` in the phonon endpoint. "
+            "Please migrate to using `get_bandstructure_from_phonon_id` with the `phonon_id` kwarg.",
+            stacklevel=2,
+            category=MPRestWarning,
+        )
+        return self.get_bandstructure_from_phonon_id(material_id, phonon_method)
+
+    def get_dos_from_phonon_id(
+        self, phonon_id: str, phonon_method: str
     ) -> PhononDOS | dict[str, Any]:
-        """Get the phonon dos pymatgen object associated with a given material ID and phonon method.
+        """Get the phonon dos pymatgen object associated with a given phonon ID and phonon method.
 
         Arguments:
-            material_id (str): Material ID for the phonon dos calculation
+            phonon_id (str): Phonon ID for the phonon dos calculation
             phonon_method (str): phonon method, i.e. pheasy or dfpt
 
         Returns:
             dos (PhononDOS): PhononDOS object
         """
         result = self._query_open_data(
             bucket="materialsproject-parsed",
-            key=f"ph-dos/{phonon_method}/{material_id}.json.gz",
+            key=f"ph-dos/{phonon_method}/{phonon_id}.json.gz",
         )[0][0]
 
         return (
@@ -115,41 +146,86 @@ def get_dos_from_material_id(
             else result  # type: ignore[return-value]
         )
 
-    def get_forceconstants_from_material_id(
-        self, material_id: str
-    ) -> list[list[Matrix3D]]:
-        """Get the force constants associated with a given material ID.
+    def get_dos_from_material_id(
+        self, material_id: str, phonon_method: str
+    ) -> PhononDOS | dict[str, Any]:
+        """Deprecated: use `get_dos_from_phonon_id` instead."""
+        warnings.warn(
+            "`material_id` has been replaced by `identifier` in the phonon endpoint. "
+            "Please migrate to using `get_dos_from_phonon_id` with the `phonon_id` kwarg.",
+            stacklevel=2,
+            category=MPRestWarning,
+        )
+        return self.get_dos_from_phonon_id(material_id, phonon_method)
+
+    def get_forceconstants_from_phonon_id(self, phonon_id: str) -> list[list[Matrix3D]]:
+        """Get the force constants associated with a given phonon ID.
 
         Arguments:
-            material_id (str): Material ID for the force constants calculation
+            phonon_id (str): Phonon ID for the force constants calculation
 
         Returns:
-            force constants (list[list[Matrix3D]]): PhononDOS object
+            force constants (list[list[Matrix3D]]): force constants
         """
         return self._query_open_data(  # type: ignore[return-value]
             bucket="materialsproject-parsed",
-            key=f"ph-force-constants/{material_id}.json.gz",
+            key=f"ph-force-constants/{phonon_id}.json.gz",
         )[0][0]
 
-    def compute_thermo_quantities(self, material_id: str, phonon_method: str):
-        """Compute thermodynamical quantities for given material ID and phonon_method.
+    def get_forceconstants_from_material_id(
+        self, material_id: str
+    ) -> list[list[Matrix3D]]:
+        """Deprecated: use `get_forceconstants_from_phonon_id` instead."""
+        warnings.warn(
+            "`material_id` has been replaced by `identifier` in the phonon endpoint. "
+            "Please migrate to using `get_forceconstants_from_phonon_id` with the `phonon_id` kwarg.",
+            stacklevel=2,
+            category=MPRestWarning,
+        )
+        return self.get_forceconstants_from_phonon_id(material_id)
+
+    def compute_thermo_quantities(
+        self,
+        phonon_id: str | None = None,
+        phonon_method: str | None = None,
+        **kwargs,
+    ):
+        """Compute thermodynamical quantities for given phonon ID and phonon_method.
 
         Arguments:
-            material_id (str): Material ID to calculate quantities for
+            phonon_id (str): Phonon ID to calculate quantities for
             phonon_method (str): phonon method, i.e. pheasy or dfpt
+            **kwargs : used for handling deprecated kwargs
 
         Returns:
             quantities (dict): thermodynamical quantities
         """
+        if "material_id" in kwargs:
+            if phonon_id:
+                raise MPRestError(
+                    "You have specified both `phonon_id` and the deprecated `material_id` tag. "
+                    "Please specify only `phonon_id`."
+                )
+            phonon_id = kwargs.pop("material_id")
+            warnings.warn(
+                "`material_id` has been replaced by `identifier` in the phonon endpoint. "
+                "Please migrate to using the newer field name and the more generic `phonon_id` kwarg.",
+                stacklevel=2,
+                category=MPRestWarning,
+            )
+
+        if phonon_id is None:
+            raise MPRestError("`phonon_id` must be specified.")
+
         use_document_model = self.use_document_model
         self.use_document_model = False
-        docs = self.search(material_ids=material_id, phonon_method=phonon_method)
+        docs = self.search(phonon_ids=phonon_id, phonon_method=phonon_method)
         if not docs or not docs[0]:
             raise MPRestError("No phonon document found")
 
         self.use_document_model = True
-        docs[0]["phonon_dos"] = self.get_dos_from_material_id(  # type: ignore[index]
-            material_id, phonon_method
+        docs[0]["phonon_dos"] = self.get_dos_from_phonon_id(  # type: ignore[index]
+            phonon_id, phonon_method
         )
         doc = PhononBSDOSDoc(**docs[0])  # type: ignore[arg-type]
         self.use_document_model = use_document_model

tests/client/materials/test_phonon.py

@@ -2,11 +2,9 @@
 
 import numpy as np
 import pytest
-
 from emmet.core.phonon import PhononBS, PhononDOS
 
 from mp_api._test_utils import client_search_testing, requires_api_key
-
 from mp_api.client.core.exceptions import MPRestError
 from mp_api.client.routes.materials.phonon import PhononRester
 
@@ -23,17 +21,22 @@ def test_phonon_search():
                 "fields",
             ],
             alt_name_dict={
-                "material_ids": "material_id",
+                "phonon_ids": "identifier",
             },
             custom_field_tests={
+                # test search kwarg backwards compat
                 "material_ids": ["mp-149", "mp-13"],
-                "material_ids": "mp-149",
+                "phonon_ids": ["ft", "mp-13"],
+                "phonon_ids": "mp-149",
                 "phonon_method": "dfpt",
             },
             sub_doc_fields=[],
         )
 
 
+# NOTE: below funcs still query legacy jsonl s3 objects
+#       they are key by 'mp-123' -> don't change id search strings
+#       to Alpha version
 @requires_api_key
 @pytest.mark.parametrize("use_document_model", [True, False])
 def test_phonon_get_methods(use_document_model):