revise lazy loading, correct rester init

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -49,6 +49,7 @@
     from typing import Any, Callable
 
     from pydantic.fields import FieldInfo
+    from mp_api.client.core.utils import LazyImport
 
 try:
     __version__ = version("mp_api")
@@ -1352,6 +1353,40 @@ def __str__(self):  # pragma: no cover
             f"Available fields: {', '.join(self.available_fields)}\n\n"
         )
 
+class CoreRester(BaseRester):
+    """Define a BaseRester with extra features for core resters.
+
+    Enables lazy importing / initialization of sub resters
+    provided in `_sub_resters`, which should be a map
+    of endpoints names to LazyImport objects.
+
+    """
+    _sub_resters : dict[str,LazyImport] = {}
+
+    def __getattr__(self, v: str):
+        if v in self._sub_resters:
+            if self._sub_resters[v]._obj is None:
+                # TODO: Enable monty decoding when tasks and SNL schema is normalized
+                monty_disable = self._sub_resters[v]._class_name in [
+                    "TaskRester",
+                    "ProvenanceRester",
+                ]
+
+                self._sub_resters[v](
+                    api_key=self.api_key,
+                    endpoint=self.endpoint.split(self.suffix)[0],
+                    include_user_agent=self._include_user_agent,
+                    session=self.session,
+                    monty_decode=False if monty_disable else self.monty_decode,
+                    use_document_model=self.use_document_model,
+                    headers=self.headers,
+                    mute_progress_bars=self.mute_progress_bars,
+                )
+            return self._sub_resters[v]
+
+    def __dir__(self):
+        return dir(self.__class__) + list(self._sub_resters)
+
 
 class MPRestError(Exception):
     """Raised when the query has problems, e.g., bad query format."""

mp_api/client/core/utils.py

@@ -76,35 +76,76 @@ def validate_ids(id_list: list[str]) -> list[str]:
 
 
 class LazyImport:
-    """Lazily import and load an object."""
+    """Lazily import and load an object.
+    
+    This class is super lazy, in that it lazily imports and caches an object.
+    If the object is a function, the function itself will be cached.
 
-    __slots__ = ["_module_name", "_class_name", "_obj"]
+    If the object is a class, and the class is initialized, the
+    current instance of the class will be cached.
 
-    def __init__(  # noqa: D107
-        self,
-        module_name: str,
-        class_name: str,
-    ) -> None:
-        self._module_name = module_name
-        self._class_name = class_name
-        self._obj = None
+    Parameters
+    -----------
+    import_str : str
+        A dot-separated, import-like string.
+    """
+
+    __slots__ = ["_module_name", "_class_name", "_obj","_imported"]
+
+    def __init__(self, import_str: str,) -> None:
+        if len(
+            splitted := import_str.rsplit(".",1)
+        ) > 1:
+            self._module_name, self._class_name = splitted
+        else:
+            self._module_name = splitted[0]
+            self._class_name = None
+
+        self._imported : Any | None = None
+        self._obj : Any | None = None
 
     def __str__(self) -> str:
-        return f"LazyImport of {self._module_name}.{self._class_name}"
+        return f"LazyImport of {self._module_name}" + (f".{self._class_name}" if self._class_name else "")
 
     def __repr__(self) -> str:
         return self.__str__()
 
-    def __call__(self, *args, **kwargs):
-        if self._obj is None:
-            try:
-                self._obj = getattr(import_module(self._module_name), self._class_name)(
-                    *args,
-                    **kwargs,
-                )
-            except Exception as exc:
-                raise ImportError(f"Failed to import {self._class_name}:\n{exc}")
-        return self._obj
-
-    def __getattr__(self, v):
-        return getattr(self._obj, v)
+    def _load(self,) -> None:
+        try:
+            _imported = import_module(self._module_name)
+            if self._class_name:
+                _imported = getattr(_imported, self._class_name)
+            self._imported = _imported
+        except Exception as exc:
+            raise ImportError(f"Failed to import {self._module_name}.{self._class_name}:\n{exc}")
+
+    def __call__(self, *args, **kwargs) -> Any:
+        """Call a function or (re-)initialize a class.
+        
+        If the object itself has not been imported, this will first import it.
+
+        If the object is a class, it will be initialized, cached, and returned.
+
+        If the object is a function, it will be cached, and this will return
+        the value(s) of the function at (*args,**kwargs).
+        """
+        if self._imported is None:
+            self._load()
+
+        if isinstance(self._imported,type):
+            self._obj = self._imported(*args, **kwargs)
+            return self._obj
+        else:
+            self._obj = self._imported
+            return self._obj(*args,**kwargs)
+
+    def __getattr__(self, v : str) -> Any:
+        """Get an attribute on a super lazy object."""
+        if self._obj is not None and hasattr(self._obj, v):
+            return getattr(self._obj, v)
+
+        if self._imported is None:
+            self._load()
+        if hasattr(self._imported,v):
+            return getattr(self._imported,v)
+        

mp_api/client/mprester.py

@@ -28,7 +28,9 @@
 from mp_api.client.core.settings import MAPIClientSettings
 from mp_api.client.core.utils import LazyImport, load_json, validate_ids
 from mp_api.client.routes import GENERIC_RESTERS
-from mp_api.client.routes.materials.materials import MATERIALS_RESTERS
+from mp_api.client.routes.materials import MATERIALS_RESTERS
+from mp_api.client.routes.molecules import MOLECULES_RESTERS
+
 
 if TYPE_CHECKING:
     from typing import Any, Literal
@@ -41,15 +43,20 @@
 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": MATERIALS_RESTERS["materials"],
     "doi": MATERIALS_RESTERS["doi"],
+    **{
+        f"molecules/{k}": v
+        for k, v in MOLECULES_RESTERS.items()
+        if k not in {"molecules",}
+    },
     **GENERIC_RESTERS,
-    "molecules/core": LazyImport("mp_api.client.routes.molecules", "MoleculeRester"),
 }
 
 
@@ -200,82 +207,10 @@ def __init__(
         }
 
         # Set remaining top level resters, or get an attribute-class name mapping
-        # for all sub-resters
-        _sub_rester_suffix_map = {"materials": {}, "molecules": {}}
-
-        # for cls in self._all_resters:
-        #     if cls.suffix not in core_suffix:
-        #         suffix_split = cls.suffix.split("/")
-
-        #         if len(suffix_split) == 1:
-        #             # Disable monty decode on nested data which may give errors
-        #             monty_disable = cls.__name__ in ["TaskRester", "ProvenanceRester"]
-        #             monty_decode = False if monty_disable else self.monty_decode
-        #             rester = cls(
-        #                 api_key=api_key,
-        #                 endpoint=self.endpoint,
-        #                 include_user_agent=include_user_agent,
-        #                 session=self.session,
-        #                 monty_decode=monty_decode,
-        #                 use_document_model=self.use_document_model,
-        #                 headers=self.headers,
-        #                 mute_progress_bars=self.mute_progress_bars,
-        #             )  # type: BaseRester
-        #             setattr(
-        #                 self,
-        #                 suffix_split[0],
-        #                 rester,
-        #             )
-        #         else:
-        #             attr = "_".join(suffix_split[1:])
-        #             if "materials" in suffix_split:
-        #                 _sub_rester_suffix_map["materials"][attr] = cls
-        #             elif "molecules" in suffix_split:
-        #                 _sub_rester_suffix_map["molecules"][attr] = cls
-
-        # TODO: Enable monty decoding when tasks and SNL schema is normalized
-        #
-        # Allow lazy loading of nested resters under materials and molecules using custom __getattr__ methods
-        def __core_custom_getattr(_self, _attr, _rester_map):
-            if _attr in RESTER_LAYOUT:
-                lazy_rester = RESTER_LAYOUT[_attr]
-                monty_disable = lazy_rester._class_name in [
-                    "TaskRester",
-                    "ProvenanceRester",
-                ]
-                monty_decode = False if monty_disable else self.monty_decode
-                rester = lazy_rester(
-                    api_key=api_key,
-                    endpoint=self.endpoint,
-                    include_user_agent=include_user_agent,
-                    session=self.session,
-                    monty_decode=monty_decode,
-                    use_document_model=self.use_document_model,
-                    headers=self.headers,
-                    mute_progress_bars=self.mute_progress_bars,
-                )  # type: BaseRester
-
-                return rester
-            else:
-                raise AttributeError(
-                    f"{_self.__class__.__name__!r} object has no attribute {_attr!r}"
-                )
-
-        def __materials_getattr__(_self, attr):
-            rester = __core_custom_getattr(_self, attr, MATERIALS_RESTERS)
-            return rester
-
-        def __molecules_getattr__(_self, attr):
-            _rester_map = _sub_rester_suffix_map["molecules"]
-            rester = __core_custom_getattr(_self, attr, _rester_map)
-            return rester
 
         for attr, rester in core_resters.items():
             setattr(self, attr, rester)
 
-        # self.materials.__getattr__ = __materials_getattr__  # type: ignore
-        # MoleculeRester.__getattr__ = __molecules_getattr__  # type: ignore
-
     @property
     def contribs(self):
         if self._contribs is None:

mp_api/client/routes/__init__.py

@@ -3,7 +3,7 @@
 from mp_api.client.core.utils import LazyImport
 
 GENERIC_RESTERS = {
-    k: LazyImport(f"mp_api.client.routes.{k}", v)
+    k: LazyImport(f"mp_api.client.routes.{k}.{v}")
     for k, v in {
         "_general_store": "GeneralStoreRester",
         "_messages": "MessagesRester",

mp_api/client/routes/materials/__init__.py

@@ -4,12 +4,11 @@
 from mp_api.client.core.utils import LazyImport
 
 MATERIALS_RESTERS: dict[str, LazyImport] = {
-    route: LazyImport(f"mp_api.client.routes.materials.{module_name}", cls_name)
-    for route, module_name, cls_name in [
+    route: LazyImport(f"mp_api.client.routes.materials.{module_name}.{cls_name}")
+    for route, module_name, cls_name in (
         ("absorption", "absorption", "AbsorptionRester"),
         ("alloys", "alloys", "AlloysRester"),
         ("bonds", "bonds", "BondsRester"),
-        ("charge_density", "charge_density", "ChargeDensityRester"),
         (
             "chemenv",
             "chemenv",
@@ -33,7 +32,7 @@
         ("materials", "materials", "MaterialsRester"),
         ("oxidation_states", "oxidation_states", "OxidationStatesRester"),
         ("phonon", "phonon", "PhononRester"),
-        ("piezoelectric", "piezoelectric", "PiezoRester"),
+        ("piezoelectric", "piezo", "PiezoRester"),
         ("provenance", "provenance", "ProvenanceRester"),
         ("robocrys", "robocrys", "RobocrysRester"),
         ("similarity", "similarity", "SimilarityRester"),
@@ -44,5 +43,5 @@
         ("tasks", "tasks", "TaskRester"),
         ("thermo", "thermo", "ThermoRester"),
         ("xas", "xas", "XASRester"),
-    ]
+    )
 }

mp_api/client/routes/materials/materials.py

@@ -5,71 +5,19 @@
 from emmet.core.vasp.material import MaterialsDoc
 from pymatgen.core.structure import Structure
 
-from mp_api.client.core import BaseRester, MPRestError
+from mp_api.client.core.client import CoreRester, MPRestError
 from mp_api.client.core.utils import validate_ids
 from mp_api.client.routes.materials import MATERIALS_RESTERS
 
 _EMMET_SETTINGS = EmmetSettings()  # type: ignore
 
 
-class MaterialsRester(BaseRester):
+class MaterialsRester(CoreRester):
     suffix = "materials/core"
     document_model = MaterialsDoc  # type: ignore
     supports_versions = True
     primary_key = "material_id"
-    _sub_resters = [
-        "eos",
-        "similarity",
-        "tasks",
-        "xas",
-        "grain_boundaries",
-        "substrates",
-        "surface_properties",
-        "phonon",
-        "elasticity",
-        "thermo",
-        "dielectric",
-        "piezoelectric",
-        "magnetism",
-        "summary",
-        "robocrys",
-        "synthesis",
-        "insertion_electrodes",
-        "conversion_electrodes",
-        "electronic_structure",
-        "electronic_structure_bandstructure",
-        "electronic_structure_dos",
-        "oxidation_states",
-        "provenance",
-        "bonds",
-        "alloys",
-        "absorption",
-        "chemenv",
-    ]
-
-    def __getattr__(self, v: str):
-        if v in self._sub_resters:
-            if MATERIALS_RESTERS[v]._obj is None:
-                # TODO: Enable monty decoding when tasks and SNL schema is normalized
-                monty_disable = MATERIALS_RESTERS[v]._class_name in [
-                    "TaskRester",
-                    "ProvenanceRester",
-                ]
-
-                MATERIALS_RESTERS[v](
-                    api_key=self.api_key,
-                    endpoint=self.endpoint.split(self.suffix)[0],
-                    include_user_agent=self._include_user_agent,
-                    session=self.session,
-                    monty_decode=False if monty_disable else self.monty_decode,
-                    use_document_model=self.use_document_model,
-                    headers=self.headers,
-                    mute_progress_bars=self.mute_progress_bars,
-                )
-            return MATERIALS_RESTERS[v]
-
-    def __dir__(self):
-        return dir(MaterialsRester) + self._sub_resters
+    _sub_resters = MATERIALS_RESTERS
 
     def get_structure_by_material_id(
         self, material_id: str, final: bool = True