ElementsQuery split from FormulaQuery

src/mp_api/bandstructure/resources.py

@@ -4,7 +4,7 @@
 from mp_api.bandstructure.models.core import BSPathType
 
 from mp_api.core.query_operator import PaginationQuery, SparseFieldsQuery, SortQuery
-from mp_api.materials.query_operators import FormulaQuery, MinMaxQuery
+from mp_api.materials.query_operators import ElementsQuery, FormulaQuery, MinMaxQuery
 from mp_api.bandstructure.query_operators import BSDataQuery
 
 from mp_api.core.utils import STORE_PARAMS
@@ -25,13 +25,10 @@ def custom_bs_endpoint_prep(self):
 
         async def get_object(
             key: str = Query(
-                ...,
-                alias=key_name,
-                title=f"The {key_name} of the {model_name} to get",
+                ..., alias=key_name, title=f"The {key_name} of the {model_name} to get",
             ),
             path_type: BSPathType = Query(
-                ...,
-                title="The k-path convention type for the band structure object",
+                ..., title="The k-path convention type for the band structure object",
             ),
             fields: STORE_PARAMS = Depends(field_input),
         ):
@@ -83,6 +80,7 @@ async def get_object(
         query_operators=[
             BSDataQuery(),
             FormulaQuery(),
+            ElementsQuery(),
             MinMaxQuery(),
             SortQuery(),
             PaginationQuery(),

src/mp_api/dos/resources.py

@@ -4,7 +4,7 @@
 from fastapi.param_functions import Query
 
 from mp_api.core.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
-from mp_api.materials.query_operators import FormulaQuery, MinMaxQuery
+from mp_api.materials.query_operators import ElementsQuery, FormulaQuery, MinMaxQuery
 from mp_api.dos.query_operators import DOSDataQuery
 
 from mp_api.core.utils import STORE_PARAMS
@@ -25,9 +25,7 @@ def custom_dos_endpoint_prep(self):
 
         async def get_object(
             key: str = Query(
-                ...,
-                alias=key_name,
-                title=f"The {key_name} of the {model_name} to get",
+                ..., alias=key_name, title=f"The {key_name} of the {model_name} to get",
             ),
             fields: STORE_PARAMS = Depends(field_input),
         ):
@@ -79,6 +77,7 @@ async def get_object(
         query_operators=[
             DOSDataQuery(),
             FormulaQuery(),
+            ElementsQuery(),
             MinMaxQuery(),
             SortQuery(),
             PaginationQuery(),

src/mp_api/materials/query_operators.py

@@ -18,6 +18,23 @@ def 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))
+
+        return {"criteria": crit}
+
+
+class ElementsQuery(QueryOperator):
+    """
+    Factory method to generate a dependency for querying by element data
+    """
+
+    def query(
+        self,
         elements: Optional[str] = Query(
             None,
             description="Query by elements in the material composition as a comma-separated list",
@@ -30,9 +47,6 @@ def query(
 
         crit = {}
 
-        if formula:
-            crit.update(formula_to_criteria(formula))
-
         if elements or exclude_elements:
             crit["elements"] = {}
 

src/mp_api/materials/resources.py

@@ -11,6 +11,7 @@
     SortQuery,
 )
 from mp_api.materials.query_operators import (
+    ElementsQuery,
     FormulaQuery,
     DeprecationQuery,
     MinMaxQuery,
@@ -252,6 +253,7 @@ async def formula_autocomplete(
         query_operators=[
             VersionQuery(),
             FormulaQuery(),
+            ElementsQuery(),
             MultiTaskIDQuery(),
             SymmetryQuery(),
             DeprecationQuery(),

src/mp_api/search/resources.py

@@ -8,6 +8,7 @@
 from mp_api.core.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
 from mp_api.core.resource import Resource
 from mp_api.materials.query_operators import (
+    ElementsQuery,
     FormulaQuery,
     MinMaxQuery,
     SymmetryQuery,
@@ -128,6 +129,7 @@ async def generate_stats(
         query_operators=[
             MultipleTaskIDsQuery(),
             FormulaQuery(),
+            ElementsQuery(),
             MinMaxQuery(),
             SymmetryQuery(),
             ThermoEnergySearchQuery(),

src/mp_api/tasks/resources.py

@@ -7,7 +7,7 @@
 
 from mp_api.core.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
 from mp_api.tasks.query_operators import MultipleTaskIDsQuery
-from mp_api.materials.query_operators import FormulaQuery
+from mp_api.materials.query_operators import ElementsQuery, FormulaQuery
 
 from monty.json import jsanitize
 import json
@@ -23,6 +23,7 @@ def task_resource(task_store):
         TaskDoc,
         query_operators=[
             FormulaQuery(),
+            ElementsQuery(),
             MultipleTaskIDsQuery(),
             SortQuery(),
             PaginationQuery(),

src/mp_api/xas/resources.py

@@ -2,7 +2,7 @@
 from mp_api.xas.models import XASDoc
 
 from mp_api.core.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
-from mp_api.materials.query_operators import FormulaQuery
+from mp_api.materials.query_operators import ElementsQuery, FormulaQuery
 from mp_api.xas.query_operator import XASQuery, XASTaskIDQuery
 
 
@@ -12,6 +12,7 @@ def xas_resource(xas_store):
         XASDoc,
         query_operators=[
             FormulaQuery(),
+            ElementsQuery(),
             XASQuery(),
             XASTaskIDQuery(),
             SortQuery(),