diff --git a/src/mp_api/electrodes/query_operators.py b/src/mp_api/electrodes/query_operators.py index d65c7df9..a983b516 100644 --- a/src/mp_api/electrodes/query_operators.py +++ b/src/mp_api/electrodes/query_operators.py @@ -1,11 +1,39 @@ from typing import Optional from fastapi import Query from mp_api.core.query_operator import STORE_PARAMS, QueryOperator +from mp_api.materials.utils import formula_to_criteria from pymatgen.core.periodic_table import Element from collections import defaultdict +class ElectrodeFormulaQuery(QueryOperator): + """ + Method to generate a query for framework formula data + """ + + def query( + self, + formula: Optional[str] = Query( + None, + description="Query by formula including anonymized formula or by including wild cards", + ), + ) -> STORE_PARAMS: + + crit = {} + + if formula: + crit.update(formula_to_criteria(formula)) + + for key in list(crit): + if "composition_reduced" in key: + framework_entry = "framework.{}".format(key.split(".")[1]) + crit[framework_entry] = crit[key] + crit.pop(key) + + return {"criteria": crit} + + class VoltageStepQuery(QueryOperator): """ Method to generate a query for ranges of voltage step data values diff --git a/src/mp_api/electrodes/resources.py b/src/mp_api/electrodes/resources.py index 32cc631b..72386f12 100644 --- a/src/mp_api/electrodes/resources.py +++ b/src/mp_api/electrodes/resources.py @@ -6,6 +6,7 @@ VoltageStepQuery, InsertionVoltageStepQuery, InsertionElectrodeQuery, + ElectrodeFormulaQuery, ) from mp_api.materials.query_operators import ElementsQuery @@ -15,6 +16,7 @@ def insertion_electrodes_resource(insertion_electrodes_store): insertion_electrodes_store, InsertionElectrodeDoc, query_operators=[ + ElectrodeFormulaQuery(), ElementsQuery(), VoltageStepQuery(), InsertionVoltageStepQuery(),