From 45468eec9bf08ea6a92a689c7ffbd77aaf12f69e Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 10:18:03 -0700 Subject: [PATCH 1/7] move to list comprehension --- mp_api/client/mprester.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py index c3fffcbb3..3fdc07f92 100644 --- a/mp_api/client/mprester.py +++ b/mp_api/client/mprester.py @@ -1203,10 +1203,11 @@ def get_entries_in_chemsys( elements_set = set(elements) # remove duplicate elements - all_chemsyses = [] - for i in range(len(elements_set)): - for els in itertools.combinations(elements_set, i + 1): - all_chemsyses.append("-".join(sorted(els))) + all_chemsyses = [ + "-".join(sorted(els)) + for i in range(len(elements_set)) + for els in itertools.combinations(elements_set, i + 1) + ] entries = [] From 2be12f8cbe8a390cf04fc1f44b92fbbb0e427b64 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 13:58:16 -0700 Subject: [PATCH 2/7] remove generic type annos --- mp_api/client/core/client.py | 15 +++------ mp_api/client/routes/_general_store.py | 2 +- mp_api/client/routes/_messages.py | 2 +- mp_api/client/routes/_user_settings.py | 2 +- mp_api/client/routes/materials/absorption.py | 2 +- mp_api/client/routes/materials/alloys.py | 2 +- mp_api/client/routes/materials/bonds.py | 2 +- mp_api/client/routes/materials/chemenv.py | 2 +- mp_api/client/routes/materials/dielectric.py | 2 +- mp_api/client/routes/materials/doi.py | 2 +- mp_api/client/routes/materials/elasticity.py | 2 +- .../routes/materials/electronic_structure.py | 2 +- mp_api/client/routes/materials/eos.py | 2 +- .../routes/materials/grain_boundaries.py | 2 +- mp_api/client/routes/materials/magnetism.py | 2 +- mp_api/client/routes/materials/materials.py | 2 +- .../routes/materials/oxidation_states.py | 2 +- mp_api/client/routes/materials/phonon.py | 2 +- mp_api/client/routes/materials/piezo.py | 2 +- mp_api/client/routes/materials/provenance.py | 2 +- mp_api/client/routes/materials/robocrys.py | 2 +- mp_api/client/routes/materials/similarity.py | 2 +- mp_api/client/routes/materials/substrates.py | 2 +- mp_api/client/routes/materials/summary.py | 2 +- .../routes/materials/surface_properties.py | 2 +- mp_api/client/routes/materials/synthesis.py | 2 +- mp_api/client/routes/materials/tasks.py | 32 ++++++++++++++++--- mp_api/client/routes/materials/thermo.py | 2 +- mp_api/client/routes/materials/xas.py | 2 +- mp_api/client/routes/molecules/bonds.py | 2 +- mp_api/client/routes/molecules/jcesr.py | 2 +- mp_api/client/routes/molecules/molecules.py | 2 +- mp_api/client/routes/molecules/orbitals.py | 2 +- .../routes/molecules/partial_charges.py | 2 +- .../client/routes/molecules/partial_spins.py | 2 +- mp_api/client/routes/molecules/redox.py | 2 +- mp_api/client/routes/molecules/summary.py | 2 +- mp_api/client/routes/molecules/tasks.py | 2 +- mp_api/client/routes/molecules/thermo.py | 2 +- mp_api/client/routes/molecules/vibrations.py | 2 +- 40 files changed, 70 insertions(+), 53 deletions(-) diff --git a/mp_api/client/core/client.py b/mp_api/client/core/client.py index 38d94bf67..e9d4be2d6 100644 --- a/mp_api/client/core/client.py +++ b/mp_api/client/core/client.py @@ -21,9 +21,7 @@ from typing import ( TYPE_CHECKING, ForwardRef, - Generic, Optional, - TypeVar, get_args, ) from urllib.parse import quote, urljoin @@ -65,10 +63,7 @@ SETTINGS = MAPIClientSettings() # type: ignore -T = TypeVar("T") - - -class BaseRester(Generic[T]): +class BaseRester: """Base client class with core stubs.""" suffix: str = "" @@ -1151,7 +1146,7 @@ def _query_resource_data( suburl: str | None = None, use_document_model: bool | None = None, timeout: int | None = None, - ) -> list[T] | list[dict]: + ) -> list[BaseModel] | list[dict]: """Query the endpoint for a list of documents without associated meta information. Only returns a single page of results. @@ -1181,7 +1176,7 @@ def _search( all_fields: bool = True, fields: list[str] | None = None, **kwargs, - ) -> list[T] | list[dict]: + ) -> list[BaseModel] | list[dict]: """A generic search method to retrieve documents matching specific parameters. Arguments: @@ -1216,7 +1211,7 @@ def get_data_by_id( self, document_id: str, fields: list[str] | None = None, - ) -> T | dict: + ) -> BaseModel | dict: warnings.warn( "get_data_by_id is deprecated and will be removed soon. Please use the search method instead.", DeprecationWarning, @@ -1251,7 +1246,7 @@ def _get_all_documents( fields=None, chunk_size=1000, num_chunks=None, - ) -> list[T] | list[dict]: + ) -> list[BaseModel] | list[dict]: """Iterates over pages until all documents are retrieved. Displays progress using tqdm. This method is designed to give a common implementation for the search_* methods on various endpoints. See diff --git a/mp_api/client/routes/_general_store.py b/mp_api/client/routes/_general_store.py index 0176c5a4f..659d06068 100644 --- a/mp_api/client/routes/_general_store.py +++ b/mp_api/client/routes/_general_store.py @@ -5,7 +5,7 @@ from mp_api.client.core import BaseRester -class GeneralStoreRester(BaseRester[GeneralStoreDoc]): # pragma: no cover +class GeneralStoreRester(BaseRester): # pragma: no cover suffix = "_general_store" document_model = GeneralStoreDoc # type: ignore primary_key = "submission_id" diff --git a/mp_api/client/routes/_messages.py b/mp_api/client/routes/_messages.py index 6526e7ca2..64f796ad0 100644 --- a/mp_api/client/routes/_messages.py +++ b/mp_api/client/routes/_messages.py @@ -7,7 +7,7 @@ from mp_api.client.core import BaseRester -class MessagesRester(BaseRester[MessagesDoc]): # pragma: no cover +class MessagesRester(BaseRester): # pragma: no cover suffix = "_messages" document_model = MessagesDoc # type: ignore primary_key = "title" diff --git a/mp_api/client/routes/_user_settings.py b/mp_api/client/routes/_user_settings.py index 90f54aaca..0f8d0bf39 100644 --- a/mp_api/client/routes/_user_settings.py +++ b/mp_api/client/routes/_user_settings.py @@ -5,7 +5,7 @@ from mp_api.client.core import BaseRester -class UserSettingsRester(BaseRester[UserSettingsDoc]): # pragma: no cover +class UserSettingsRester(BaseRester): # pragma: no cover suffix = "_user_settings" document_model = UserSettingsDoc # type: ignore primary_key = "consumer_id" diff --git a/mp_api/client/routes/materials/absorption.py b/mp_api/client/routes/materials/absorption.py index 9ef7ff338..c8d2c0c45 100644 --- a/mp_api/client/routes/materials/absorption.py +++ b/mp_api/client/routes/materials/absorption.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class AbsorptionRester(BaseRester[AbsorptionDoc]): +class AbsorptionRester(BaseRester): suffix = "materials/absorption" document_model = AbsorptionDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/alloys.py b/mp_api/client/routes/materials/alloys.py index 53b79064a..fc07d2a58 100644 --- a/mp_api/client/routes/materials/alloys.py +++ b/mp_api/client/routes/materials/alloys.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class AlloysRester(BaseRester[AlloyPairDoc]): +class AlloysRester(BaseRester): suffix = "materials/alloys" document_model = AlloyPairDoc # type: ignore primary_key = "pair_id" diff --git a/mp_api/client/routes/materials/bonds.py b/mp_api/client/routes/materials/bonds.py index bdbe4ca2d..9009bc4b5 100644 --- a/mp_api/client/routes/materials/bonds.py +++ b/mp_api/client/routes/materials/bonds.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class BondsRester(BaseRester[BondingDoc]): +class BondsRester(BaseRester): suffix = "materials/bonds" document_model = BondingDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/chemenv.py b/mp_api/client/routes/materials/chemenv.py index bc05fff85..d3a029fd7 100644 --- a/mp_api/client/routes/materials/chemenv.py +++ b/mp_api/client/routes/materials/chemenv.py @@ -14,7 +14,7 @@ from mp_api.client.core.utils import validate_ids -class ChemenvRester(BaseRester[ChemEnvDoc]): +class ChemenvRester(BaseRester): suffix = "materials/chemenv" document_model = ChemEnvDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/dielectric.py b/mp_api/client/routes/materials/dielectric.py index 4adedc37e..2723f94cd 100644 --- a/mp_api/client/routes/materials/dielectric.py +++ b/mp_api/client/routes/materials/dielectric.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class DielectricRester(BaseRester[DielectricDoc]): +class DielectricRester(BaseRester): suffix = "materials/dielectric" document_model = DielectricDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/doi.py b/mp_api/client/routes/materials/doi.py index bb66f185d..bc8e766b6 100644 --- a/mp_api/client/routes/materials/doi.py +++ b/mp_api/client/routes/materials/doi.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class DOIRester(BaseRester[DOIDoc]): +class DOIRester(BaseRester): suffix = "doi" document_model = DOIDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/elasticity.py b/mp_api/client/routes/materials/elasticity.py index e6c416968..d2099cdca 100644 --- a/mp_api/client/routes/materials/elasticity.py +++ b/mp_api/client/routes/materials/elasticity.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class ElasticityRester(BaseRester[ElasticityDoc]): +class ElasticityRester(BaseRester): suffix = "materials/elasticity" document_model = ElasticityDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/electronic_structure.py b/mp_api/client/routes/materials/electronic_structure.py index c43a88afd..ccaaaad83 100644 --- a/mp_api/client/routes/materials/electronic_structure.py +++ b/mp_api/client/routes/materials/electronic_structure.py @@ -16,7 +16,7 @@ from mp_api.client.core.utils import validate_ids -class ElectronicStructureRester(BaseRester[ElectronicStructureDoc]): +class ElectronicStructureRester(BaseRester): suffix = "materials/electronic_structure" document_model = ElectronicStructureDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/eos.py b/mp_api/client/routes/materials/eos.py index d77dfd459..26a14fef2 100644 --- a/mp_api/client/routes/materials/eos.py +++ b/mp_api/client/routes/materials/eos.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class EOSRester(BaseRester[EOSDoc]): +class EOSRester(BaseRester): suffix = "materials/eos" document_model = EOSDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/grain_boundaries.py b/mp_api/client/routes/materials/grain_boundaries.py index 1721ccd64..fa11bbfac 100644 --- a/mp_api/client/routes/materials/grain_boundaries.py +++ b/mp_api/client/routes/materials/grain_boundaries.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class GrainBoundaryRester(BaseRester[GrainBoundaryDoc]): +class GrainBoundaryRester(BaseRester): suffix = "materials/grain_boundaries" document_model = GrainBoundaryDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/magnetism.py b/mp_api/client/routes/materials/magnetism.py index 96c3c59a6..da19f7274 100644 --- a/mp_api/client/routes/materials/magnetism.py +++ b/mp_api/client/routes/materials/magnetism.py @@ -9,7 +9,7 @@ from mp_api.client.core.utils import validate_ids -class MagnetismRester(BaseRester[MagnetismDoc]): +class MagnetismRester(BaseRester): suffix = "materials/magnetism" document_model = MagnetismDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/materials.py b/mp_api/client/routes/materials/materials.py index a7bff5c46..df24e93b4 100644 --- a/mp_api/client/routes/materials/materials.py +++ b/mp_api/client/routes/materials/materials.py @@ -40,7 +40,7 @@ _EMMET_SETTINGS = EmmetSettings() # type: ignore -class MaterialsRester(BaseRester[MaterialsDoc]): +class MaterialsRester(BaseRester): suffix = "materials/core" document_model = MaterialsDoc # type: ignore supports_versions = True diff --git a/mp_api/client/routes/materials/oxidation_states.py b/mp_api/client/routes/materials/oxidation_states.py index fca6e93de..769b4ef81 100644 --- a/mp_api/client/routes/materials/oxidation_states.py +++ b/mp_api/client/routes/materials/oxidation_states.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class OxidationStatesRester(BaseRester[OxidationStateDoc]): +class OxidationStatesRester(BaseRester): suffix = "materials/oxidation_states" document_model = OxidationStateDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/phonon.py b/mp_api/client/routes/materials/phonon.py index 605591048..24eb17ec3 100644 --- a/mp_api/client/routes/materials/phonon.py +++ b/mp_api/client/routes/materials/phonon.py @@ -9,7 +9,7 @@ from mp_api.client.core.utils import validate_ids -class PhononRester(BaseRester[PhononBSDOSDoc]): +class PhononRester(BaseRester): suffix = "materials/phonon" document_model = PhononBSDOSDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/piezo.py b/mp_api/client/routes/materials/piezo.py index 306cf5069..ee76e1cd6 100644 --- a/mp_api/client/routes/materials/piezo.py +++ b/mp_api/client/routes/materials/piezo.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class PiezoRester(BaseRester[PiezoelectricDoc]): +class PiezoRester(BaseRester): suffix = "materials/piezoelectric" document_model = PiezoelectricDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/provenance.py b/mp_api/client/routes/materials/provenance.py index b4ff240fc..da837e17c 100644 --- a/mp_api/client/routes/materials/provenance.py +++ b/mp_api/client/routes/materials/provenance.py @@ -6,7 +6,7 @@ from mp_api.client.core.utils import validate_ids -class ProvenanceRester(BaseRester[ProvenanceDoc]): +class ProvenanceRester(BaseRester): suffix = "materials/provenance" document_model = ProvenanceDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/robocrys.py b/mp_api/client/routes/materials/robocrys.py index bf24ccc65..e4a15a41c 100644 --- a/mp_api/client/routes/materials/robocrys.py +++ b/mp_api/client/routes/materials/robocrys.py @@ -6,7 +6,7 @@ from mp_api.client.core.utils import validate_ids -class RobocrysRester(BaseRester[RobocrystallogapherDoc]): +class RobocrysRester(BaseRester): suffix = "materials/robocrys" document_model = RobocrystallogapherDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/similarity.py b/mp_api/client/routes/materials/similarity.py index c2ae9e999..cf1dc175e 100644 --- a/mp_api/client/routes/materials/similarity.py +++ b/mp_api/client/routes/materials/similarity.py @@ -6,7 +6,7 @@ from mp_api.client.core.utils import validate_ids -class SimilarityRester(BaseRester[SimilarityDoc]): +class SimilarityRester(BaseRester): suffix = "materials/similarity" document_model = SimilarityDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/substrates.py b/mp_api/client/routes/materials/substrates.py index ac2ba6c8e..6c3f90d14 100644 --- a/mp_api/client/routes/materials/substrates.py +++ b/mp_api/client/routes/materials/substrates.py @@ -7,7 +7,7 @@ from mp_api.client.core import BaseRester -class SubstratesRester(BaseRester[SubstratesDoc]): +class SubstratesRester(BaseRester): suffix = "materials/substrates" document_model = SubstratesDoc # type: ignore primary_key = "film_id" diff --git a/mp_api/client/routes/materials/summary.py b/mp_api/client/routes/materials/summary.py index e04d71f3e..b90782d08 100644 --- a/mp_api/client/routes/materials/summary.py +++ b/mp_api/client/routes/materials/summary.py @@ -11,7 +11,7 @@ from mp_api.client.core.utils import validate_ids -class SummaryRester(BaseRester[SummaryDoc]): +class SummaryRester(BaseRester): suffix = "materials/summary" document_model = SummaryDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/surface_properties.py b/mp_api/client/routes/materials/surface_properties.py index a3c92c247..8b9d53f63 100644 --- a/mp_api/client/routes/materials/surface_properties.py +++ b/mp_api/client/routes/materials/surface_properties.py @@ -8,7 +8,7 @@ from mp_api.client.core.utils import validate_ids -class SurfacePropertiesRester(BaseRester[SurfacePropDoc]): +class SurfacePropertiesRester(BaseRester): suffix = "materials/surface_properties" document_model = SurfacePropDoc # type: ignore primary_key = "material_id" diff --git a/mp_api/client/routes/materials/synthesis.py b/mp_api/client/routes/materials/synthesis.py index 171f5e2b4..6788814ce 100644 --- a/mp_api/client/routes/materials/synthesis.py +++ b/mp_api/client/routes/materials/synthesis.py @@ -9,7 +9,7 @@ from mp_api.client.core import BaseRester, MPRestError -class SynthesisRester(BaseRester[SynthesisSearchResultModel]): +class SynthesisRester(BaseRester): suffix = "materials/synthesis" document_model = SynthesisSearchResultModel # type: ignore diff --git a/mp_api/client/routes/materials/tasks.py b/mp_api/client/routes/materials/tasks.py index 5e12f9aa1..4f06d8260 100644 --- a/mp_api/client/routes/materials/tasks.py +++ b/mp_api/client/routes/materials/tasks.py @@ -1,17 +1,20 @@ from __future__ import annotations from datetime import datetime +from typing import TYPE_CHECKING from emmet.core.tasks import TaskDoc from mp_api.client.core import BaseRester, MPRestError from mp_api.client.core.utils import validate_ids +if TYPE_CHECKING: + from pydantic import BaseModel -class TaskRester(BaseRester[TaskDoc]): - suffix = "materials/tasks" - document_model = TaskDoc # type: ignore - primary_key = "task_id" +class TaskRester(BaseRester): + suffix : str = "materials/tasks" + document_model : type[BaseModel] = TaskDoc # type: ignore + primary_key : str = "task_id" def get_trajectory(self, task_id): """Returns a Trajectory object containing the geometry of the @@ -32,6 +35,21 @@ def get_trajectory(self, task_id): raise MPRestError(f"No trajectory data for {task_id} found") return traj_data + + def _set_entry(self, task) -> dict | TaskDoc: + """Get the ComputedEntry corresponding to a task.""" + cr = task.calcs_reversed if hasattr(task,"calcs_reversed") else task.get("calcs_reversed") + task_id = task.task_id if hasattr(task,"task_id") else task.get("task_id") + if not cr or not task_id: + return None + entry = TaskDoc.get_entry(cr,task_id) + if not self.monty_decode: + entry = entry.as_dict() + if isinstance(task,dict): + task["entry"] = entry + else: + task.entry = entry + return task def search( self, @@ -89,10 +107,14 @@ def search( } ) - return super()._search( + tasks = super()._search( num_chunks=num_chunks, chunk_size=chunk_size, all_fields=all_fields, fields=fields, **query_params, ) + + # if not fields or "entry" in fields: + # tasks = [self._set_entry(task) for task in tasks] + return tasks \ No newline at end of file diff --git a/mp_api/client/routes/materials/thermo.py b/mp_api/client/routes/materials/thermo.py index 509745efa..edf06dd8d 100644 --- a/mp_api/client/routes/materials/thermo.py +++ b/mp_api/client/routes/materials/thermo.py @@ -16,7 +16,7 @@ from emmet.core.thermo import ThermoType -class ThermoRester(BaseRester[ThermoDoc]): +class ThermoRester(BaseRester): suffix = "materials/thermo" document_model = ThermoDoc # type: ignore supports_versions = True diff --git a/mp_api/client/routes/materials/xas.py b/mp_api/client/routes/materials/xas.py index 1b2e28a60..99dae9c9e 100644 --- a/mp_api/client/routes/materials/xas.py +++ b/mp_api/client/routes/materials/xas.py @@ -22,7 +22,7 @@ ) -class XASRester(BaseRester[XASDoc]): +class XASRester(BaseRester): suffix = "materials/xas" document_model = XASDoc # type: ignore primary_key = "spectrum_id" diff --git a/mp_api/client/routes/molecules/bonds.py b/mp_api/client/routes/molecules/bonds.py index 3ea6d79e7..96a1d0d1c 100644 --- a/mp_api/client/routes/molecules/bonds.py +++ b/mp_api/client/routes/molecules/bonds.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesBondRester(BaseRester[MoleculeBondingDoc]): +class MoleculesBondRester(BaseRester): suffix = "molecules/bonding" document_model = MoleculeBondingDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/jcesr.py b/mp_api/client/routes/molecules/jcesr.py index b8bdb2b4b..15216e44b 100644 --- a/mp_api/client/routes/molecules/jcesr.py +++ b/mp_api/client/routes/molecules/jcesr.py @@ -9,7 +9,7 @@ from mp_api.client.core.utils import validate_ids -class JcesrMoleculesRester(BaseRester[MoleculesDoc]): +class JcesrMoleculesRester(BaseRester): suffix = "molecules/jcesr" document_model = MoleculesDoc # type: ignore primary_key = "task_id" diff --git a/mp_api/client/routes/molecules/molecules.py b/mp_api/client/routes/molecules/molecules.py index c278980ed..922f136ea 100644 --- a/mp_api/client/routes/molecules/molecules.py +++ b/mp_api/client/routes/molecules/molecules.py @@ -13,7 +13,7 @@ _EMMET_SETTINGS = EmmetSettings() -class BaseMoleculeRester(BaseRester[MoleculeDoc]): +class BaseMoleculeRester(BaseRester): document_model = MoleculeDoc primary_key = "molecule_id" diff --git a/mp_api/client/routes/molecules/orbitals.py b/mp_api/client/routes/molecules/orbitals.py index 123b27d9d..9b93cab09 100644 --- a/mp_api/client/routes/molecules/orbitals.py +++ b/mp_api/client/routes/molecules/orbitals.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesOrbitalsRester(BaseRester[OrbitalDoc]): +class MoleculesOrbitalsRester(BaseRester): suffix = "molecules/orbitals" document_model = OrbitalDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/partial_charges.py b/mp_api/client/routes/molecules/partial_charges.py index 8f0d99b86..1c828c500 100644 --- a/mp_api/client/routes/molecules/partial_charges.py +++ b/mp_api/client/routes/molecules/partial_charges.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesPartialChargesRester(BaseRester[PartialChargesDoc]): +class MoleculesPartialChargesRester(BaseRester): suffix = "molecules/partial_charges" document_model = PartialChargesDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/partial_spins.py b/mp_api/client/routes/molecules/partial_spins.py index da290db38..a04504a5e 100644 --- a/mp_api/client/routes/molecules/partial_spins.py +++ b/mp_api/client/routes/molecules/partial_spins.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesPartialSpinsRester(BaseRester[PartialSpinsDoc]): +class MoleculesPartialSpinsRester(BaseRester): suffix = "molecules/partial_spins" document_model = PartialSpinsDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/redox.py b/mp_api/client/routes/molecules/redox.py index 8d6ebd5da..e41c2cda1 100644 --- a/mp_api/client/routes/molecules/redox.py +++ b/mp_api/client/routes/molecules/redox.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesRedoxRester(BaseRester[RedoxDoc]): +class MoleculesRedoxRester(BaseRester): suffix = "molecules/redox" document_model = RedoxDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/summary.py b/mp_api/client/routes/molecules/summary.py index 6e340476a..7cce3a974 100644 --- a/mp_api/client/routes/molecules/summary.py +++ b/mp_api/client/routes/molecules/summary.py @@ -8,7 +8,7 @@ from mp_api.client.core import BaseRester -class MoleculesSummaryRester(BaseRester[MoleculeSummaryDoc]): +class MoleculesSummaryRester(BaseRester): suffix = "molecules/summary" document_model = MoleculeSummaryDoc # type: ignore primary_key = "molecule_id" diff --git a/mp_api/client/routes/molecules/tasks.py b/mp_api/client/routes/molecules/tasks.py index 0969b3a05..02289aaca 100644 --- a/mp_api/client/routes/molecules/tasks.py +++ b/mp_api/client/routes/molecules/tasks.py @@ -6,7 +6,7 @@ from mp_api.client.core.utils import validate_ids -class MoleculesTaskRester(BaseRester[TaskDocument]): +class MoleculesTaskRester(BaseRester): suffix = "molecules/tasks" document_model = TaskDocument primary_key = "task_id" diff --git a/mp_api/client/routes/molecules/thermo.py b/mp_api/client/routes/molecules/thermo.py index 6a2c59c2b..d2e1b0c1b 100644 --- a/mp_api/client/routes/molecules/thermo.py +++ b/mp_api/client/routes/molecules/thermo.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesThermoRester(BaseRester[MoleculeThermoDoc]): +class MoleculesThermoRester(BaseRester): suffix = "molecules/thermo" document_model = MoleculeThermoDoc primary_key = "property_id" diff --git a/mp_api/client/routes/molecules/vibrations.py b/mp_api/client/routes/molecules/vibrations.py index 6d6d1245a..146936a9c 100644 --- a/mp_api/client/routes/molecules/vibrations.py +++ b/mp_api/client/routes/molecules/vibrations.py @@ -6,7 +6,7 @@ from mp_api.client.core import BaseRester -class MoleculesVibrationRester(BaseRester[VibrationDoc]): +class MoleculesVibrationRester(BaseRester): suffix = "molecules/vibrations" document_model = VibrationDoc primary_key = "property_id" From 23c5beb35daf301d6c3be28363aed405c5a3eec0 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 14:23:42 -0700 Subject: [PATCH 3/7] pydantic 2.12.3 fix --- mp_api/client/core/client.py | 7 ++++--- mp_api/client/routes/materials/tasks.py | 23 ++++++++++++++--------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/mp_api/client/core/client.py b/mp_api/client/core/client.py index e9d4be2d6..9fb66d28e 100644 --- a/mp_api/client/core/client.py +++ b/mp_api/client/core/client.py @@ -63,6 +63,7 @@ SETTINGS = MAPIClientSettings() # type: ignore + class BaseRester: """Base client class with core stubs.""" @@ -1072,7 +1073,9 @@ def _generate_returned_model( include_fields: dict[str, tuple[type, FieldInfo]] = {} for name in set_fields: field_copy = model_fields[name]._copy() - field_copy.default = None + if not getattr(field_copy,"default_factory"): + # Fields with a default_factory cannot also have a default in pydantic>=2.12.3 + field_copy.default = None include_fields[name] = ( Optional[model_fields[name].annotation], field_copy, @@ -1092,8 +1095,6 @@ def _generate_returned_model( ), __module__=self.document_model.__module__, ) - # if other_vars: - # data_model.model_rebuild(_types_namespace=other_vars) orig_rester_name = self.document_model.__name__ diff --git a/mp_api/client/routes/materials/tasks.py b/mp_api/client/routes/materials/tasks.py index 4f06d8260..08cdb53ca 100644 --- a/mp_api/client/routes/materials/tasks.py +++ b/mp_api/client/routes/materials/tasks.py @@ -11,10 +11,11 @@ if TYPE_CHECKING: from pydantic import BaseModel + class TaskRester(BaseRester): - suffix : str = "materials/tasks" - document_model : type[BaseModel] = TaskDoc # type: ignore - primary_key : str = "task_id" + suffix: str = "materials/tasks" + document_model: type[BaseModel] = TaskDoc # type: ignore + primary_key: str = "task_id" def get_trajectory(self, task_id): """Returns a Trajectory object containing the geometry of the @@ -35,17 +36,21 @@ def get_trajectory(self, task_id): raise MPRestError(f"No trajectory data for {task_id} found") return traj_data - + def _set_entry(self, task) -> dict | TaskDoc: """Get the ComputedEntry corresponding to a task.""" - cr = task.calcs_reversed if hasattr(task,"calcs_reversed") else task.get("calcs_reversed") - task_id = task.task_id if hasattr(task,"task_id") else task.get("task_id") + cr = ( + task.calcs_reversed + if hasattr(task, "calcs_reversed") + else task.get("calcs_reversed") + ) + task_id = task.task_id if hasattr(task, "task_id") else task.get("task_id") if not cr or not task_id: return None - entry = TaskDoc.get_entry(cr,task_id) + entry = TaskDoc.get_entry(cr, task_id) if not self.monty_decode: entry = entry.as_dict() - if isinstance(task,dict): + if isinstance(task, dict): task["entry"] = entry else: task.entry = entry @@ -117,4 +122,4 @@ def search( # if not fields or "entry" in fields: # tasks = [self._set_entry(task) for task in tasks] - return tasks \ No newline at end of file + return tasks From fdf1745f2e72276b44c444ff16630fc37c6dae67 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 14:26:22 -0700 Subject: [PATCH 4/7] cleanup --- mp_api/client/core/client.py | 13 +++-------- mp_api/client/routes/materials/tasks.py | 29 +++---------------------- 2 files changed, 6 insertions(+), 36 deletions(-) diff --git a/mp_api/client/core/client.py b/mp_api/client/core/client.py index 9fb66d28e..a195c6c2f 100644 --- a/mp_api/client/core/client.py +++ b/mp_api/client/core/client.py @@ -136,15 +136,8 @@ def __init__( if not self.endpoint.endswith("/"): self.endpoint += "/" - if session: - self._session = session - else: - self._session = None # type: ignore - - if s3_client: - self._s3_client = s3_client - else: - self._s3_client = None + self._session = session + self._s3_client = s3_client @property def session(self) -> requests.Session: @@ -592,7 +585,7 @@ def _submit_requests( # noqa url: url used to make request use_document_model: if None, will defer to the self.use_document_model attribute parallel_param: parameter to parallelize requests with - num_chu: fieldsnky: Maximum number of chunks of data to yield. None will yield all possible. + num_chunks: Maximum number of chunks of data to yield. None will yield all possible. chunk_size: Number of data entries per chunk. timeout: Time in seconds to wait until a request timeout error is thrown diff --git a/mp_api/client/routes/materials/tasks.py b/mp_api/client/routes/materials/tasks.py index 08cdb53ca..0b55329eb 100644 --- a/mp_api/client/routes/materials/tasks.py +++ b/mp_api/client/routes/materials/tasks.py @@ -36,26 +36,7 @@ def get_trajectory(self, task_id): raise MPRestError(f"No trajectory data for {task_id} found") return traj_data - - def _set_entry(self, task) -> dict | TaskDoc: - """Get the ComputedEntry corresponding to a task.""" - cr = ( - task.calcs_reversed - if hasattr(task, "calcs_reversed") - else task.get("calcs_reversed") - ) - task_id = task.task_id if hasattr(task, "task_id") else task.get("task_id") - if not cr or not task_id: - return None - entry = TaskDoc.get_entry(cr, task_id) - if not self.monty_decode: - entry = entry.as_dict() - if isinstance(task, dict): - task["entry"] = entry - else: - task.entry = entry - return task - + def search( self, task_ids: str | list[str] | None = None, @@ -112,14 +93,10 @@ def search( } ) - tasks = super()._search( + return super()._search( num_chunks=num_chunks, chunk_size=chunk_size, all_fields=all_fields, fields=fields, **query_params, - ) - - # if not fields or "entry" in fields: - # tasks = [self._set_entry(task) for task in tasks] - return tasks + ) \ No newline at end of file From ed40ecbc5db2ec73b247a15588ea8f4971360110 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 14:26:54 -0700 Subject: [PATCH 5/7] use core task doc --- mp_api/client/routes/materials/tasks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mp_api/client/routes/materials/tasks.py b/mp_api/client/routes/materials/tasks.py index 0b55329eb..7cb4c0078 100644 --- a/mp_api/client/routes/materials/tasks.py +++ b/mp_api/client/routes/materials/tasks.py @@ -3,7 +3,7 @@ from datetime import datetime from typing import TYPE_CHECKING -from emmet.core.tasks import TaskDoc +from emmet.core.tasks import CoreTaskDoc from mp_api.client.core import BaseRester, MPRestError from mp_api.client.core.utils import validate_ids @@ -14,7 +14,7 @@ class TaskRester(BaseRester): suffix: str = "materials/tasks" - document_model: type[BaseModel] = TaskDoc # type: ignore + document_model: type[BaseModel] = CoreTaskDoc # type: ignore primary_key: str = "task_id" def get_trajectory(self, task_id): From 14d71c22d1bd897db09d5e0ad66f1985d8b29197 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 14:27:41 -0700 Subject: [PATCH 6/7] precommit --- mp_api/client/core/client.py | 2 +- mp_api/client/routes/materials/tasks.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mp_api/client/core/client.py b/mp_api/client/core/client.py index a195c6c2f..4cd98958c 100644 --- a/mp_api/client/core/client.py +++ b/mp_api/client/core/client.py @@ -1066,7 +1066,7 @@ def _generate_returned_model( include_fields: dict[str, tuple[type, FieldInfo]] = {} for name in set_fields: field_copy = model_fields[name]._copy() - if not getattr(field_copy,"default_factory"): + if not field_copy.default_factory: # Fields with a default_factory cannot also have a default in pydantic>=2.12.3 field_copy.default = None include_fields[name] = ( diff --git a/mp_api/client/routes/materials/tasks.py b/mp_api/client/routes/materials/tasks.py index 7cb4c0078..c78650780 100644 --- a/mp_api/client/routes/materials/tasks.py +++ b/mp_api/client/routes/materials/tasks.py @@ -36,7 +36,7 @@ def get_trajectory(self, task_id): raise MPRestError(f"No trajectory data for {task_id} found") return traj_data - + def search( self, task_ids: str | list[str] | None = None, @@ -48,7 +48,7 @@ def search( chunk_size: int = 1000, all_fields: bool = True, fields: list[str] | None = None, - ) -> list[TaskDoc] | list[dict]: + ) -> list[CoreTaskDoc] | list[dict]: """Query core task docs using a variety of search criteria. Arguments: @@ -62,11 +62,11 @@ def search( 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. Max size is 100. all_fields (bool): Whether to return all fields in the document. Defaults to True. - fields (List[str]): List of fields in TaskDoc to return data for. + fields (List[str]): List of fields in CoreTaskDoc to return data for. Default is material_id, last_updated, and formula_pretty if all_fields is False. Returns: - ([TaskDoc], [dict]) List of task documents or dictionaries. + ([CoreTaskDoc], [dict]) List of task documents or dictionaries. """ query_params = {} # type: dict @@ -99,4 +99,4 @@ def search( all_fields=all_fields, fields=fields, **query_params, - ) \ No newline at end of file + ) From de750bc1c186d6654628392cf31f0bf3d71f9d6d Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Oct 2025 14:31:50 -0700 Subject: [PATCH 7/7] test newer pydantic --- requirements/requirements-ubuntu-latest_py3.11.txt | 2 +- requirements/requirements-ubuntu-latest_py3.11_extras.txt | 2 +- requirements/requirements-ubuntu-latest_py3.12.txt | 2 +- requirements/requirements-ubuntu-latest_py3.12_extras.txt | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/requirements/requirements-ubuntu-latest_py3.11.txt b/requirements/requirements-ubuntu-latest_py3.11.txt index 0dd0b0eaf..cda326c08 100644 --- a/requirements/requirements-ubuntu-latest_py3.11.txt +++ b/requirements/requirements-ubuntu-latest_py3.11.txt @@ -83,7 +83,7 @@ plotly==6.3.1 # via pymatgen pybtex==0.25.1 # via emmet-core -pydantic==2.12.2 +pydantic==2.12.3 # via # emmet-core # pydantic-settings diff --git a/requirements/requirements-ubuntu-latest_py3.11_extras.txt b/requirements/requirements-ubuntu-latest_py3.11_extras.txt index 9fddc73be..dae1e91a1 100644 --- a/requirements/requirements-ubuntu-latest_py3.11_extras.txt +++ b/requirements/requirements-ubuntu-latest_py3.11_extras.txt @@ -280,7 +280,7 @@ pycodestyle==2.14.0 # via # flake8 # mp-api (pyproject.toml) -pydantic==2.12.2 +pydantic==2.12.3 # via # emmet-core # pydantic-settings diff --git a/requirements/requirements-ubuntu-latest_py3.12.txt b/requirements/requirements-ubuntu-latest_py3.12.txt index eb743970e..f12a31abd 100644 --- a/requirements/requirements-ubuntu-latest_py3.12.txt +++ b/requirements/requirements-ubuntu-latest_py3.12.txt @@ -83,7 +83,7 @@ plotly==6.3.1 # via pymatgen pybtex==0.25.1 # via emmet-core -pydantic==2.12.2 +pydantic==2.12.3 # via # emmet-core # pydantic-settings diff --git a/requirements/requirements-ubuntu-latest_py3.12_extras.txt b/requirements/requirements-ubuntu-latest_py3.12_extras.txt index 8863d5357..be9196f2b 100644 --- a/requirements/requirements-ubuntu-latest_py3.12_extras.txt +++ b/requirements/requirements-ubuntu-latest_py3.12_extras.txt @@ -280,7 +280,7 @@ pycodestyle==2.14.0 # via # flake8 # mp-api (pyproject.toml) -pydantic==2.12.2 +pydantic==2.12.3 # via # emmet-core # pydantic-settings