New molecules resters (#803)

* Init commit of new resters

* Add tests and legacy

* Add inits and fix mprester imports

* Commit new molecules resters

* Linting

* Fix test imports

mp_api/client/mprester.py

@@ -24,6 +24,11 @@
 from mp_api.client.core.settings import MAPIClientSettings
 from mp_api.client.core.utils import validate_ids
 from mp_api.client.routes import (
+    GeneralStoreRester,
+    UserSettingsRester,
+)
+from mp_api.client.routes.legacy import LegacyMoleculesRester
+from mp_api.client.routes.materials import (
     AbsorptionRester,
     AlloysRester,
     BandStructureRester,
@@ -38,7 +43,6 @@
     ElectronicStructureRester,
     EOSRester,
     FermiRester,
-    GeneralStoreRester,
     GrainBoundaryRester,
     MagnetismRester,
     MaterialsRester,
@@ -55,9 +59,21 @@
     SynthesisRester,
     TaskRester,
     ThermoRester,
-    UserSettingsRester,
     XASRester,
 )
+from mp_api.client.routes.molecules import (
+    AssociatedMoleculeRester,
+    MoleculeRester,
+    MoleculesBondRester,
+    MoleculesOrbitalsRester,
+    MoleculesPartialChargesRester,
+    MoleculesPartialSpinsRester,
+    MoleculesRedoxRester,
+    MoleculesSummaryRester,
+    MoleculesTaskRester,
+    MoleculesThermoRester,
+    MoleculesVibrationRester,
+)
 
 _DEPRECATION_WARNING = (
     "MPRester is being modernized. Please use the new method suggested and "
@@ -79,6 +95,8 @@ class MPRester:
     # To re-generate this list, use:
     # for rester in MPRester()._all_resters:
     #     print(f"{rester.suffix.replace('/', '_')}: {rester.__class__.__name__}")
+
+    # Materials
     eos: EOSRester
     materials: MaterialsRester
     similarity: SimilarityRester
@@ -92,7 +110,6 @@ class MPRester:
     elasticity: ElasticityRester
     thermo: ThermoRester
     dielectric: DielectricRester
-    doi: DOIRester
     piezoelectric: PiezoRester
     magnetism: MagnetismRester
     summary: SummaryRester
@@ -110,6 +127,25 @@ class MPRester:
     alloys: AlloysRester
     absorption: AbsorptionRester
     chemenv: ChemenvRester
+
+    # Molecules
+    molecules_bonding: MoleculesBondRester
+    molecules_associated: AssociatedMoleculeRester
+    molecules: MoleculeRester
+    molecules_orbital: MoleculesOrbitalsRester
+    molecules_partial_charges: MoleculesPartialChargesRester
+    molecules_partial_spins: MoleculesPartialSpinsRester
+    molecules_redox: MoleculesRedoxRester
+    molecules_summary: MoleculesSummaryRester
+    molecules_tasks: MoleculesTaskRester
+    molecules_thermo: MoleculesThermoRester
+    molecules_vibrations: MoleculesVibrationRester
+
+    # Legacy
+    legacy_jcesr: LegacyMoleculesRester
+
+    # Generic
+    doi: DOIRester
     _user_settings: UserSettingsRester
     _general_store: GeneralStoreRester
 
@@ -207,39 +243,49 @@ def __init__(
         if not self.endpoint.endswith("/"):
             self.endpoint += "/"
 
-        for cls in BaseRester.__subclasses__():
-            rester = cls(
-                api_key=api_key,
-                endpoint=endpoint,
-                include_user_agent=include_user_agent,
-                session=self.session,
-                monty_decode=monty_decode
-                if cls not in [TaskRester, ProvenanceRester]  # type: ignore
-                else False,  # Disable monty decode on nested data which may give errors
-                use_document_model=use_document_model,
-                headers=self.headers,
-            )  # type: BaseRester
-
-            self._all_resters.append(rester)
-
-            suffix_split = cls.suffix.split("/")
-
-            att_map = {"legacy/jcesr": "molecules", "materials/core": "materials"}
-
-            if len(suffix_split) == 1:
-                setattr(
-                    self,
-                    cls.suffix.split("/")[0],
-                    rester,
-                )
-            else:
-                setattr(
-                    self,
-                    att_map[cls.suffix]
-                    if cls.suffix in att_map
-                    else "_".join(suffix_split[1:]),
-                    rester,
-                )
+        for _cls in BaseRester.__subclasses__():
+            sub_resters = _cls.__subclasses__()
+
+            resters = sub_resters if sub_resters else [_cls]
+
+            for cls in resters:
+                rester = cls(
+                    api_key=api_key,
+                    endpoint=endpoint,
+                    include_user_agent=include_user_agent,
+                    session=self.session,
+                    monty_decode=monty_decode
+                    if cls not in [TaskRester, ProvenanceRester]  # type: ignore
+                    else False,  # Disable monty decode on nested data which may give errors
+                    use_document_model=use_document_model,
+                    headers=self.headers,
+                )  # type: BaseRester
+
+                self._all_resters.append(rester)
+
+                suffix_split = cls.suffix.split("/")
+
+                att_map = {"molecules/core": "molecules", "materials/core": "materials"}
+
+                if len(suffix_split) == 1:
+                    setattr(
+                        self,
+                        suffix_split[0],
+                        rester,
+                    )
+                else:
+                    if cls.suffix in att_map:
+                        attr = att_map[cls.suffix]
+                    elif "materials" in suffix_split:
+                        attr = "_".join(suffix_split[1:])
+                    else:
+                        attr = "_".join(suffix_split)
+
+                    setattr(
+                        self,
+                        attr,
+                        rester,
+                    )
 
     def __enter__(self):
         """Support for "with" context."""

mp_api/client/routes/__init__.py

@@ -1,45 +1,2 @@
-from ast import Import
-
 from ._general_store import GeneralStoreRester
 from ._user_settings import UserSettingsRester
-from .absorption import AbsorptionRester
-from .bonds import BondsRester
-from .chemenv import ChemenvRester
-from .dielectric import DielectricRester
-from .doi import DOIRester
-from .elasticity import ElasticityRester
-from .electrodes import ElectrodeRester
-from .electronic_structure import (
-    BandStructureRester,
-    DosRester,
-    ElectronicStructureRester,
-)
-from .eos import EOSRester
-from .fermi import FermiRester
-from .grain_boundary import GrainBoundaryRester
-from .magnetism import MagnetismRester
-from .materials import MaterialsRester
-from .molecules import MoleculesRester
-from .oxidation_states import OxidationStatesRester
-from .phonon import PhononRester
-from .piezo import PiezoRester
-from .provenance import ProvenanceRester
-from .robocrys import RobocrysRester
-from .similarity import SimilarityRester
-from .substrates import SubstratesRester
-from .summary import SummaryRester
-from .surface_properties import SurfacePropertiesRester
-from .synthesis import SynthesisRester
-from .tasks import TaskRester
-from .thermo import ThermoRester
-from .xas import XASRester
-
-try:
-    from .alloys import AlloysRester
-except ImportError:
-    AlloysRester = None  # type: ignore
-
-try:
-    from .charge_density import ChargeDensityRester
-except ImportError:
-    ChargeDensityRester = None  # type: ignore

mp_api/client/routes/legacy/__init__.py

@@ -0,0 +1 @@
+from .jcesr import LegacyMoleculesRester

mp_api/client/routes/legacy/jcesr.py

@@ -0,0 +1,95 @@
+from collections import defaultdict
+from typing import List, Optional, Tuple
+
+from emmet.core.molecules_jcesr import MoleculesDoc
+from pymatgen.core.periodic_table import Element
+
+from mp_api.client.core import BaseRester
+
+
+class LegacyMoleculesRester(BaseRester[MoleculesDoc]):
+    suffix = "legacy/jcesr"
+    document_model = MoleculesDoc  # type: ignore
+    primary_key = "task_id"
+
+    def search(
+        self,
+        charge: Optional[Tuple[float, float]] = None,
+        elements: Optional[List[Element]] = None,
+        EA: Optional[Tuple[float, float]] = None,
+        IE: Optional[Tuple[float, float]] = None,
+        nelements: Optional[Tuple[float, float]] = None,
+        pointgroup: Optional[str] = None,
+        smiles: Optional[str] = None,
+        sort_fields: Optional[List[str]] = None,
+        num_chunks: Optional[int] = None,
+        chunk_size: int = 1000,
+        all_fields: bool = True,
+        fields: Optional[List[str]] = None,
+    ):
+        """
+        Query equations of state docs using a variety of search criteria.
+
+        Arguments:
+            charge (Tuple[float,float]): Minimum and maximum value of the charge in +e to consider.
+            elements (List[Element]): A list of elements.
+            film_orientation (List[Elements]): List of elements that are in the molecule.
+            EA (Tuple[float,float]): Minimum and maximum value of the electron affinity in eV to consider.
+            IE (Tuple[float,float]): Minimum and maximum value of the ionization energy in eV to consider.
+            nelements (Tuple[float,float]): Minimum and maximum number of elements in the molecule to consider.
+            pointgroup (str): Point group of the molecule in Schoenflies notation.
+            smiles (str): The simplified molecular input line-entry system (SMILES) representation of the molecule.
+            sort_fields (List[str]): Fields used to sort results. Prefix with '-' to sort in descending order.
+            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 MoleculesDoc to return data for.
+                Default is the material_id only if all_fields is False.
+
+        Returns:
+            ([MoleculesDoc]) List of molecule documents
+        """
+
+        query_params = defaultdict(dict)  # type: dict
+
+        if elements:
+            query_params.update({"elements": ",".join([str(ele) for ele in elements])})
+
+        if pointgroup:
+            query_params.update({"pointgroup": pointgroup})
+
+        if smiles:
+            query_params.update({"smiles": smiles})
+
+        if nelements:
+            query_params.update(
+                {"nelements_min": nelements[0], "nelements_max": nelements[1]}
+            )
+
+        if EA:
+            query_params.update({"EA_min": EA[0], "EA_max": EA[1]})
+
+        if IE:
+            query_params.update({"IE_min": IE[0], "IE_max": IE[1]})
+
+        if charge:
+            query_params.update({"charge_min": charge[0], "charge_max": charge[1]})
+
+        if sort_fields:
+            query_params.update(
+                {"_sort_fields": ",".join([s.strip() for s in sort_fields])}
+            )
+
+        query_params = {
+            entry: query_params[entry]
+            for entry in query_params
+            if query_params[entry] is not None
+        }
+
+        return super()._search(
+            num_chunks=num_chunks,
+            chunk_size=chunk_size,
+            all_fields=all_fields,
+            fields=fields,
+            **query_params
+        )

mp_api/client/routes/materials/__init__.py

@@ -0,0 +1,41 @@
+from .absorption import AbsorptionRester
+from .bonds import BondsRester
+from .chemenv import ChemenvRester
+from .dielectric import DielectricRester
+from .doi import DOIRester
+from .elasticity import ElasticityRester
+from .electrodes import ElectrodeRester
+from .electronic_structure import (
+    BandStructureRester,
+    DosRester,
+    ElectronicStructureRester,
+)
+from .eos import EOSRester
+from .fermi import FermiRester
+from .grain_boundary import GrainBoundaryRester
+from .magnetism import MagnetismRester
+from .materials import MaterialsRester
+from .molecules import MoleculesRester
+from .oxidation_states import OxidationStatesRester
+from .phonon import PhononRester
+from .piezo import PiezoRester
+from .provenance import ProvenanceRester
+from .robocrys import RobocrysRester
+from .similarity import SimilarityRester
+from .substrates import SubstratesRester
+from .summary import SummaryRester
+from .surface_properties import SurfacePropertiesRester
+from .synthesis import SynthesisRester
+from .tasks import TaskRester
+from .thermo import ThermoRester
+from .xas import XASRester
+
+try:
+    from .alloys import AlloysRester
+except ImportError:
+    AlloysRester = None  # type: ignore
+
+try:
+    from .charge_density import ChargeDensityRester
+except ImportError:
+    ChargeDensityRester = None  # type: ignore

mp_api/client/routes/molecules/__init__.py

@@ -0,0 +1,10 @@
+from .bonds import MoleculesBondRester
+from .molecules import AssociatedMoleculeRester, MoleculeRester
+from .orbitals import MoleculesOrbitalsRester
+from .partial_charges import MoleculesPartialChargesRester
+from .partial_spins import MoleculesPartialSpinsRester
+from .redox import MoleculesRedoxRester
+from .summary import MoleculesSummaryRester
+from .tasks import MoleculesTaskRester
+from .thermo import MoleculesThermoRester
+from .vibrations import MoleculesVibrationRester