diff --git a/src/mp_api/routes/search/query_operators.py b/src/mp_api/routes/search/query_operators.py index 71e725e9..357c3c8c 100644 --- a/src/mp_api/routes/search/query_operators.py +++ b/src/mp_api/routes/search/query_operators.py @@ -141,6 +141,38 @@ def ensure_indexes(self): # pragma: no cover return [("theoretical", False)] +class SearchESQuery(QueryOperator): + """ + Method to generate a query on search electronic structure data. + """ + + def query( + self, + is_gap_direct: Optional[bool] = Query( + None, description="Whether a band gap is direct or not." + ), + is_metal: Optional[bool] = Query( + None, description="Whether the material is considered a metal." + ), + ) -> STORE_PARAMS: + + crit = defaultdict(dict) # type: dict + + if is_gap_direct is not None: + crit["is_gap_direct"] = is_gap_direct + + if is_metal is not None: + crit["is_metal"] = is_metal + + return {"criteria": crit} + + def ensure_indexes(self): # pragma: no cover + + keys = ["is_gap_direct", "is_metal"] + + return [(key, False) for key in keys] + + class SearchStatsQuery(QueryOperator): """ Method to generate a query on search stats data diff --git a/src/mp_api/routes/search/resources.py b/src/mp_api/routes/search/resources.py index eb5bea60..1a0418fa 100644 --- a/src/mp_api/routes/search/resources.py +++ b/src/mp_api/routes/search/resources.py @@ -21,6 +21,7 @@ SearchIsTheoreticalQuery, SearchMagneticQuery, SearchStatsQuery, + SearchESQuery, ) @@ -36,6 +37,7 @@ def search_resource(search_store): SearchIsStableQuery(), SearchIsTheoreticalQuery(), SearchMagneticQuery(), + SearchESQuery(), NumericQuery(model=SearchDoc, excluded_fields=["composition"]), HasPropsQuery(), DeprecationQuery(),