Conversion electrode + trajectory tweaks (#1041)

Author: esoteric-ephemera

mp_api/client/routes/materials/electrodes.py

@@ -173,4 +173,5 @@ class ConversionElectrodeRester(BaseElectrodeRester):
         "elements",
         "stability_charge",
         "stability_discharge",
+        "exclude_elements",
     ]

mp_api/client/routes/materials/tasks.py

@@ -3,12 +3,15 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
+from emmet.core.mpid import MPID, AlphaID
 from emmet.core.tasks import CoreTaskDoc
 
 from mp_api.client.core import BaseRester, MPRestError
 from mp_api.client.core.utils import validate_ids
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pydantic import BaseModel
 
 
@@ -17,17 +20,21 @@ class TaskRester(BaseRester):
     document_model: type[BaseModel] = CoreTaskDoc  # type: ignore
     primary_key: str = "task_id"
 
-    def get_trajectory(self, task_id):
+    def get_trajectory(self, task_id: MPID | AlphaID | str) -> list[dict[str, Any]]:
         """Returns a Trajectory object containing the geometry of the
         material throughout a calculation. This is most useful for
         observing how a material relaxes during a geometry optimization.
 
         Args:
-            task_id (str): Task ID
+            task_id (str, MPID, AlphaID): Task ID
 
+        Returns:
+            list of dict representing emmet.core.trajectory.Trajectory
         """
         traj_data = self._query_resource_data(
-            {"task_ids": [task_id]}, suburl="trajectory/", use_document_model=False
+            {"task_ids": [AlphaID(task_id).string]},
+            suburl="trajectory/",
+            use_document_model=False,
         )[0].get(
             "trajectories", None
         )  # type: ignore

tests/materials/test_electrodes.py

@@ -77,7 +77,6 @@ def test_conversion_client(conversion_rester):
         custom_field_tests={
             "battery_ids": ["mp-1067_Al"],
             "working_ion": Element("Li"),
-            "exclude_elements": ["Co", "O"],
         },
         sub_doc_fields=sub_doc_fields,
     )

tests/materials/test_tasks.py

@@ -2,6 +2,8 @@
 from core_function import client_search_testing
 import pytest
 
+from emmet.core.mpid import MPID, AlphaID
+from emmet.core.trajectory import Trajectory
 from emmet.core.utils import utcnow
 from mp_api.client.routes.materials.tasks import TaskRester
 
@@ -52,8 +54,13 @@ def test_client(rester):
     )
 
 
-def test_get_trajectories(rester):
-    trajectories = [traj for traj in rester.get_trajectory("mp-149")]
+@pytest.mark.parametrize("mpid", ["mp-149", MPID("mp-149"), AlphaID("mp-149")])
+def test_get_trajectories(rester, mpid):
+    trajectories = [traj for traj in rester.get_trajectory(mpid)]
 
-    for traj in trajectories:
-        assert ("@module", "pymatgen.core.trajectory") in traj.items()
+    expected_model_fields = {
+        field_name
+        for field_name, field in Trajectory.model_fields.items()
+        if not field.exclude
+    }
+    assert all(set(traj) == expected_model_fields for traj in trajectories)