Swap dyanmic query op to pydantic 2

materialsproject · Sep 7, 2023 · 3b4df02 · 3b4df02
1 parent b29c6d1
commit 3b4df02
Showing 1 changed file with 37 additions and 11 deletions.
diff --git a/src/maggma/api/query_operator/dynamic.py b/src/maggma/api/query_operator/dynamic.py
@@ -5,7 +5,7 @@
 from fastapi.params import Query
 from monty.json import MontyDecoder
 from pydantic import BaseModel
-from pydantic.fields import ModelField
+from pydantic.fields import FieldInfo
 
 from maggma.api.query_operator import QueryOperator
 from maggma.api.utils import STORE_PARAMS
@@ -25,8 +25,10 @@ def __init__(
         self.fields = fields
         self.excluded_fields = excluded_fields
 
-        all_fields: Dict[str, ModelField] = model.__fields__
-        param_fields = fields or list(set(all_fields.keys()) - set(excluded_fields or []))
+        all_fields: Dict[str, FieldInfo] = model.model_fields
+        param_fields = fields or list(
+            set(all_fields.keys()) - set(excluded_fields or [])
+        )
 
         # Convert the fields into operator tuples
         ops = [
@@ -47,7 +49,9 @@ def query(**kwargs) -> STORE_PARAMS:
                     try:
                         criteria.append(self.mapping[k](v))
                     except KeyError:
-                        raise KeyError(f"Cannot find key {k} in current query to database mapping")
+                        raise KeyError(
+                            f"Cannot find key {k} in current query to database mapping"
+                        )
 
             final_crit = {}
             for entry in criteria:
@@ -78,7 +82,9 @@ def query(self):
         "Stub query function for abstract class"
 
     @abstractmethod
-    def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
+    def field_to_operator(
+        self, name: str, field: ModelField
+    ) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
         """
         Converts a PyDantic ModelField into a Tuple with the
             - query param name,
@@ -107,7 +113,9 @@ def as_dict(self) -> Dict:
 class NumericQuery(DynamicQueryOperator):
     "Query Operator to enable searching on numeric fields"
 
-    def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
+    def field_to_operator(
+        self, name: str, field: ModelField
+    ) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
         """
         Converts a PyDantic ModelField into a Tuple with the
         query_param name,
@@ -171,7 +179,11 @@ def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any
                             default=None,
                             description=f"Query for {title} being any of these values. Provide a comma separated list.",
                         ),
-                        lambda val: {f"{field.name}": {"$in": [int(entry.strip()) for entry in val.split(",")]}},
+                        lambda val: {
+                            f"{field.name}": {
+                                "$in": [int(entry.strip()) for entry in val.split(",")]
+                            }
+                        },
                     ),
                     (
                         f"{field.name}_neq_any",
@@ -181,7 +193,11 @@ def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any
                             description=f"Query for {title} being not any of these values. \
                             Provide a comma separated list.",
                         ),
-                        lambda val: {f"{field.name}": {"$nin": [int(entry.strip()) for entry in val.split(",")]}},
+                        lambda val: {
+                            f"{field.name}": {
+                                "$nin": [int(entry.strip()) for entry in val.split(",")]
+                            }
+                        },
                     ),
                 ]
             )
@@ -192,7 +208,9 @@ def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any
 class StringQueryOperator(DynamicQueryOperator):
     "Query Operator to enable searching on numeric fields"
 
-    def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
+    def field_to_operator(
+        self, name: str, field: ModelField
+    ) -> List[Tuple[str, Any, Query, Callable[..., Dict]]]:
         """
         Converts a PyDantic ModelField into a Tuple with the
         query_param name,
@@ -233,7 +251,11 @@ def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any
                         default=None,
                         description=f"Query for {title} being any of these values. Provide a comma separated list.",
                     ),
-                    lambda val: {f"{field.name}": {"$in": [entry.strip() for entry in val.split(",")]}},
+                    lambda val: {
+                        f"{field.name}": {
+                            "$in": [entry.strip() for entry in val.split(",")]
+                        }
+                    },
                 ),
                 (
                     f"{field.name}_neq_any",
@@ -242,7 +264,11 @@ def field_to_operator(self, name: str, field: ModelField) -> List[Tuple[str, Any
                         default=None,
                         description=f"Query for {title} being not any of these values. Provide a comma separated list",
                     ),
-                    lambda val: {f"{field.name}": {"$nin": [entry.strip() for entry in val.split(",")]}},
+                    lambda val: {
+                        f"{field.name}": {
+                            "$nin": [entry.strip() for entry in val.split(",")]
+                        }
+                    },
                 ),
             ]