more housekeeping

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -461,7 +461,6 @@ def _query_resource(
             url = validate_endpoint(self.endpoint, suffix=suburl)
 
             if query_s3:
-                db_version = self.db_version.replace(".", "-")
                 if "/" not in self.suffix:
                     suffix = self.suffix
                 elif self.suffix == "molecules/summary":
@@ -477,7 +476,7 @@ def _query_resource(
                     bucket_suffix, prefix = "parsed", "tasks_atomate2"
                 else:
                     bucket_suffix = "build"
-                    prefix = f"collections/{db_version}/{suffix}"
+                    prefix = f"collections/{self.db_version.replace('.', '-')}/{suffix}"
 
                 bucket = f"materialsproject-{bucket_suffix}"
                 paginator = self.s3_client.get_paginator("list_objects_v2")

mp_api/client/mprester.py

@@ -8,7 +8,6 @@
 
 from emmet.core.electronic_structure import BSPathType
 from emmet.core.mpid import MPID, AlphaID
-from emmet.core.tasks import TaskDoc
 from emmet.core.types.enums import ThermoType
 from emmet.core.vasp.calc_types import CalcType
 from packaging import version
@@ -38,30 +37,23 @@
 if TYPE_CHECKING:
     from typing import Any, Literal
 
+    from emmet.core.tasks import CoreTaskDoc
     from pymatgen.analysis.phase_diagram import PDEntry
     from pymatgen.entries.computed_entries import ComputedEntry
 
+
 DEFAULT_THERMOTYPE_CRITERIA = {"thermo_types": ["GGA_GGA+U"]}
 
 RESTER_LAYOUT = {
     "molecules/core": LazyImport(
         "mp_api.client.routes.molecules.molecules.MoleculeRester"
     ),
-    "materials/core": MATERIALS_RESTERS["materials"],
-    **{
-        f"materials/{k}": v
-        for k, v in MATERIALS_RESTERS.items()
-        if k not in {"materials", "doi"}
-    },
+    "materials/core": LazyImport(
+        "mp_api.client.routes.materials.materials.MaterialsRester"
+    ),
+    **{f"materials/{k}": v for k, v in MATERIALS_RESTERS.items() if k not in {"doi"}},
     "doi": MATERIALS_RESTERS["doi"],
-    **{
-        f"molecules/{k}": v
-        for k, v in MOLECULES_RESTERS.items()
-        if k
-        not in {
-            "molecules",
-        }
-    },
+    **{f"molecules/{k}": v for k, v in MOLECULES_RESTERS.items()},
     **GENERIC_RESTERS,
 }
 
@@ -1134,22 +1126,21 @@ def get_wulff_shape(self, material_id: str):
 
     def get_charge_density_from_task_id(
         self, task_id: str, inc_task_doc: bool = False
-    ) -> Chgcar | tuple[Chgcar, TaskDoc | dict] | None:
+    ) -> Chgcar | tuple[Chgcar, CoreTaskDoc | dict] | None:
         """Get charge density data for a given task_id.
 
         Arguments:
             task_id (str): A task id
             inc_task_doc (bool): Whether to include the task document in the returned data.
 
         Returns:
-            (Chgcar, (Chgcar, TaskDoc | dict), None): Pymatgen Chgcar object, or tuple with object and TaskDoc
+            (Chgcar, (Chgcar, CoreTaskDoc | dict), None): Pymatgen Chgcar object, or tuple with object and CoreTaskDoc
         """
-        kwargs = dict(
+        chgcar = self.materials.tasks._query_open_data(
             bucket="materialsproject-parsed",
             key=f"chgcars/{validate_ids([task_id])[0]}.json.gz",
             decoder=lambda x: load_json(x, deser=True),
-        )
-        chgcar = self.materials.tasks._query_open_data(**kwargs)[0][0]["data"]
+        )[0][0]["data"]
 
         if inc_task_doc:
             task_doc = self.materials.tasks.search(task_ids=task_id)[0]
@@ -1159,15 +1150,16 @@ def get_charge_density_from_task_id(
 
     def get_charge_density_from_material_id(
         self, material_id: str, inc_task_doc: bool = False
-    ) -> Chgcar | tuple[Chgcar, TaskDoc | dict] | None:
+    ) -> Chgcar | tuple[Chgcar, CoreTaskDoc | dict] | None:
         """Get charge density data for a given Materials Project ID.
 
         Arguments:
             material_id (str): Material Project ID
             inc_task_doc (bool): Whether to include the task document in the returned data.
 
         Returns:
-            (Chgcar, (Chgcar, TaskDoc | dict), None): Pymatgen Chgcar object, or tuple with object and TaskDoc
+            (Chgcar, (Chgcar, CoreTaskDoc | dict), None): Pymatgen Chgcar object,
+            or tuple with object and CoreTaskDoc
         """
         # TODO: really we want a recommended task_id for charge densities here
         # this could potentially introduce an ambiguity
@@ -1177,7 +1169,7 @@ def get_charge_density_from_material_id(
         if not task_ids:
             return None
 
-        results: list[TaskDoc] = self.materials.tasks.search(
+        results: list[CoreTaskDoc] = self.materials.tasks.search(
             task_ids=task_ids, fields=["last_updated", "task_id"]
         )  # type: ignore
 

mp_api/client/routes/materials/__init__.py

@@ -1,4 +1,4 @@
-"""Define routes and imports to materials resters."""
+"""Define routes and imports to non-core materials resters."""
 from __future__ import annotations
 
 from mp_api.client.core.utils import LazyImport
@@ -29,7 +29,6 @@
         ("grain_boundaries", "grain_boundaries", "GrainBoundaryRester"),
         ("insertion_electrodes", "electrodes", "ElectrodeRester"),
         ("magnetism", "magnetism", "MagnetismRester"),
-        ("materials", "materials", "MaterialsRester"),
         ("oxidation_states", "oxidation_states", "OxidationStatesRester"),
         ("phonon", "phonon", "PhononRester"),
         ("piezoelectric", "piezo", "PiezoRester"),

mp_api/client/routes/molecules/__init__.py

@@ -1,11 +1,11 @@
+"""Define routes to non-core molecules resters."""
 from __future__ import annotations
 
 from mp_api.client.core.utils import LazyImport
 
-MOLECULES_RESTERS = {
+MOLECULES_RESTERS: dict[str, LazyImport] = {
     k: LazyImport(f"mp_api.client.routes.molecules.{k}.{v}")
     for k, v in (
-        ("molecules", "MoleculeRester"),
         ("jcesr", "JcesrMoleculesRester"),
         ("summary", "MoleculesSummaryRester"),
     )

mp_api/client/routes/molecules/molecules.py

@@ -1,3 +1,10 @@
+"""Define core molecules functionality.
+
+Note that these classes are not currently working.
+The `MoleculeRester` methods: `search`, `find_molecule`, `get_molecule_by_mpculeid`,
+all return 404s when attempting to use them.
+"""
+
 from __future__ import annotations
 
 from emmet.core.mpid import MPculeID
@@ -190,6 +197,3 @@ class AssociatedMoleculeRester(BaseMoleculeRester):
 class MoleculeRester(BaseMoleculeRester):
     suffix = "molecules/core"
     _sub_resters = MOLECULES_RESTERS
-
-    def __dir__(self):
-        return dir(MoleculeRester) + self._sub_resters

tests/core/test_utils.py

@@ -7,6 +7,23 @@
 from mp_api.client.core.utils import LazyImport
 
 
+def test_lazy_import_module():
+    import_str = "typing"
+    lazy_mod = LazyImport(import_str)
+    assert lazy_mod._module_name == import_str
+    assert lazy_mod._class_name is None
+    assert lazy_mod.__repr__() == str(lazy_mod)
+    lazy_any = lazy_mod.Any
+
+    from typing import Any
+
+    assert lazy_any == Any
+
+    # Ensure failure to load raises error
+    with pytest.raises(ImportError, match="Failed to import"):
+        _ = LazyImport("tieping")._load()
+
+
 def test_lazy_import_function():
     import_str = "json.dumps"
     lazy_func = LazyImport(import_str)

tests/materials/test_materials.py

@@ -5,6 +5,7 @@
 from emmet.core.symmetry import CrystalSystem
 
 from mp_api.client.routes.materials.materials import MaterialsRester
+from mp_api.client.routes.materials import MATERIALS_RESTERS
 
 
 @pytest.fixture
@@ -55,6 +56,10 @@ def rester():
 def test_client(rester):
     search_method = rester.search
 
+    assert sorted(dir(rester)) == sorted(
+        dir(rester.__class__) + list(MATERIALS_RESTERS)
+    )
+
     client_search_testing(
         search_method=search_method,
         excluded_params=excluded_params,

tests/molecules/test_molecules.py

@@ -0,0 +1,10 @@
+"""Test basic molecules functionality."""
+
+from mp_api.client.routes.molecules.molecules import BaseMoleculeRester, MoleculeRester
+from mp_api.client.routes.molecules import MOLECULES_RESTERS
+
+
+def test_molecule_rester():
+    assert set(dir(MoleculeRester())) == set(
+        dir(BaseMoleculeRester) + list(MOLECULES_RESTERS)
+    )

tests/test_mprester.py

@@ -365,6 +365,8 @@ def test_get_charge_density_from_material_id(self, mpr):
         assert isinstance(chgcar, Chgcar)
         assert isinstance(TaskDoc.model_validate(task_doc.model_dump()), TaskDoc)
 
+        assert mpr.get_charge_density_from_material_id("mp-0") is None
+
     def test_get_charge_density_from_task_id(self, mpr):
         chgcar = mpr.get_charge_density_from_task_id("mp-2246557")
         assert isinstance(chgcar, Chgcar)