reorg tests slightly

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -191,7 +191,7 @@ def __init__(
             warnings.warn(
                 "Ignoring `monty_decode`, as it is no longer a supported option in `mp_api`."
                 "The client by default returns results consistent with `monty_decode=True`.",
-                category=DeprecationWarning,
+                category=MPRestWarning,
                 stacklevel=2,
             )
 
@@ -1349,7 +1349,7 @@ def new_str(self) -> str:
 
             return (
                 f"\033[4m\033[1m{self.__class__.__name__}"
-                f"<{self.__class__.__base__.__name__}>\033[0;0m\033[0;0m"
+                f"<{orig_rester_name}>\033[0;0m\033[0;0m"
                 f"\n{extra}\n\n"
                 f"\033[1mFields not requested:\033[0;0m\n{fields_not_requested}"
             )
@@ -1597,7 +1597,7 @@ def __getattr__(self, v: str):
                 self.sub_resters[v](
                     api_key=self.api_key,
                     endpoint=self.base_endpoint,
-                    include_user_agent=self._include_user_agent,
+                    include_user_agent=self.include_user_agent,
                     session=self.session,
                     use_document_model=self.use_document_model,
                     headers=self.headers,
@@ -1606,6 +1606,7 @@ def __getattr__(self, v: str):
                     force_renew=self.force_renew,
                 )
             return self.sub_resters[v]
+        raise AttributeError(f"{self.__class__} has no attribute {v}")
 
     def __dir__(self):
         return dir(self.__class__) + list(self._sub_resters)

mp_api/client/mprester.py

@@ -1076,25 +1076,28 @@ def get_entries_in_chemsys(
         if isinstance(elements, str):
             elements = elements.split("-")
 
-        elements_set = set(elements)  # remove duplicate elements
+        # 9 elements would be sum_{i=1}^{9} (9 choose i) = 511
+        # From testing, this is the highest number of chemsys
+        # we can query before URI lengths are exceeded
+        if len(elements_set := set(elements)) > 9:  # remove duplicate elements
+            raise MPRestError(
+                "Please specify fewer elements to query by, "
+                "or identify a subset of relevant chemical systems to query first."
+            )
 
         all_chemsyses = [
             "-".join(sorted(els))
             for i in range(len(elements_set))
             for els in itertools.combinations(elements_set, i + 1)
         ]
 
-        entries = []
-
-        entries.extend(
-            self.get_entries(
-                all_chemsyses,
-                compatible_only=compatible_only,
-                property_data=property_data,
-                conventional_unit_cell=conventional_unit_cell,
-                additional_criteria=additional_criteria or DEFAULT_THERMOTYPE_CRITERIA,
-                **kwargs,
-            )
+        entries = self.get_entries(
+            all_chemsyses,
+            compatible_only=compatible_only,
+            property_data=property_data,
+            conventional_unit_cell=conventional_unit_cell,
+            additional_criteria=additional_criteria or DEFAULT_THERMOTYPE_CRITERIA,
+            **kwargs,
         )
 
         if use_gibbs:
@@ -1261,21 +1264,41 @@ def get_charge_density_from_material_id(
         task_id = latest_doc["task_id"]
         return self.get_charge_density_from_task_id(task_id, inc_task_doc)
 
-    def get_download_info(self, material_ids, calc_types=None, file_patterns=None):
+    def get_download_info(
+        self,
+        material_ids: str | MPID | list[str | MPID],
+        calc_types: list[str | CalcType] | None = None,
+        file_patterns: list[str] | None = None,
+    ):
         """Get a list of URLs to retrieve raw VASP output files from the NoMaD repository
         Args:
-            material_ids (list): list of material identifiers (mp-id's)
-            task_types (list): list of task types to include in download (see CalcType Enum class)
+            material_ids (str or MPID, or list thereof): list of material identifiers (mp-id's)
+            calc_types (list of str or CalcType): list of calc types to include in download (see CalcType Enum class)
             file_patterns (list): list of wildcard file names to include for each task
         Returns:
             a tuple of 1) a dictionary mapping material_ids to task_ids and
             calc_types, and 2) a list of URLs to download zip archives from
             NoMaD repository. Each zip archive will contain a manifest.json with
             metadata info, e.g. the task/external_ids that belong to a directory.
         """
+        warnings.warn(
+            "Full downloads of raw data are being transitioned to "
+            "Materials Project's AWS S3 OpenData buckets. "
+            "These features for accessing legacy raw data via NOMAD "
+            "are maintained but may not be supported in the future.",
+            category=MPRestWarning,
+            stacklevel=2,
+        )
+
         # task_id's correspond to NoMaD external_id's
+        if isinstance(material_ids, str | MPID):
+            material_ids = [material_ids]
+
         calc_types = (
-            [t.value for t in calc_types if isinstance(t, CalcType)]
+            [
+                t.value if isinstance(t, CalcType) else CalcType(t).value
+                for t in calc_types
+            ]
             if calc_types
             else []
         )

mp_api/client/routes/materials/materials.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 from emmet.core.symmetry import CrystalSystem
@@ -172,11 +173,11 @@ def search(
 
     def find_structure(
         self,
-        filename_or_structure,
+        filename_or_structure: str | Path | Structure,
         ltol=MAPI_CLIENT_SETTINGS.LTOL,
         stol=MAPI_CLIENT_SETTINGS.STOL,
         angle_tol=MAPI_CLIENT_SETTINGS.ANGLE_TOL,
-        allow_multiple_results=False,
+        allow_multiple_results: bool | int = False,
     ) -> list[str] | str:
         """Finds matching structures from the Materials Project database.
 
@@ -186,48 +187,75 @@ def find_structure(
         default tolerances.
 
         Args:
-            filename_or_structure: filename or Structure object
+            filename_or_structure: filename as a str or Path, or a Structure object
             ltol: fractional length tolerance
             stol: site tolerance
             angle_tol: angle tolerance in degrees
-            allow_multiple_results: changes return type for either
-            a single material_id or list of material_ids
+            allow_multiple_results (bool or int): changes return type for either
+                a single material_id or list of material_ids.
+                If a bool, returns either all matches (True) or one match at most (False).
+                If an int, returns that many matches at most.
+
         Returns:
             A matching material_id if one is found or list of results if allow_multiple_results
             is True
         Raises:
             MPRestError
         """
-        params = {"ltol": ltol, "stol": stol, "angle_tol": angle_tol, "_limit": 1}
+        from pymatgen.analysis.structure_matcher import (
+            ElementComparator,
+            StructureMatcher,
+        )
 
-        if isinstance(filename_or_structure, str):
+        if (
+            isinstance(filename_or_structure, str | Path)
+            and Path(filename_or_structure).exists()
+        ):
             s = Structure.from_file(filename_or_structure)
         elif isinstance(filename_or_structure, Structure):
             s = filename_or_structure
         else:
             raise MPRestError("Provide filename or Structure object.")
 
-        results = self._post_resource(
-            body=s.as_dict(),
-            params=params,
-            suburl="find_structure",
-            use_document_model=False,
-        ).get("data")
-
-        if not results:
+        mat_docs = self.search(
+            formula=s.reduced_formula, fields=["material_id", "structure"]
+        )
+        if not mat_docs:
             return []
 
-        material_ids = validate_ids([doc["material_id"] for doc in results])
+        if isinstance(allow_multiple_results, bool):
+            max_matches: int = len(mat_docs) if allow_multiple_results else 1
+        elif isinstance(allow_multiple_results, int):
+            max_matches = allow_multiple_results
+        else:
+            raise MPRestError(
+                f"`allow_multiple_results` must be a bool or int, not {type(allow_multiple_results)}"
+            )
 
-        if len(material_ids) > 1:  # type: ignore
-            if not allow_multiple_results:
-                raise ValueError(
-                    "Multiple matches found for this combination of tolerances, but "
-                    "`allow_multiple_results` set to False."
-                )
-            return material_ids  # type: ignore
+        matcher = StructureMatcher(
+            ltol=ltol,
+            stol=stol,
+            angle_tol=angle_tol,
+            primitive_cell=True,
+            scale=True,
+            attempt_supercell=False,
+            comparator=ElementComparator(),
+        )
 
-        return material_ids[0]
+        matches: list[str] = []
+        for doc in mat_docs:
+            if matcher.fit(
+                s,
+                doc.structure if self.use_document_model else Structure.from_dict(doc["structure"]),  # type: ignore
+            ):
+                matches.append(doc.material_id.string if self.use_document_model else doc["material_id"])  # type: ignore
+                if len(matches) >= max_matches:
+                    break
+
+        if not matches:
+            return []
+        material_ids = validate_ids(matches)
+        return material_ids if allow_multiple_results else material_ids[0]
 
     def get_blessed_entries(
         self,