Merge pull request #221 from Materials-Consortia/ml-evs/not_and_queries

Handle arbitrary nested NOT/AND/OR in queries
Materials-Consortia · Mar 13, 2020 · 5569bda · 5569bda
2 parents 2ed7f74 + ee1a400
commit 5569bda
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 10 deletions.
diff --git a/optimade/filtertransformers/mongo.py b/optimade/filtertransformers/mongo.py
@@ -88,16 +88,7 @@ def expression_clause(self, arg):
 
     def expression_phrase(self, arg):
         # expression_phrase: [ NOT ] ( comparison | "(" expression ")" )
-        if len(arg) == 1:
-            # without NOT
-            return arg[0]
-
-        if list(arg[1].keys()) == ["$or"]:
-            return {"$nor": arg[1]["$or"]}
-
-        # with NOT
-        # TODO: This implementation probably fails in the case of `"(" expression ")"`
-        return {prop: {"$not": expr} for prop, expr in arg[1].items()}
+        return self._recursive_expression_phrase(arg)
 
     @v_args(inline=True)
     def comparison(self, value):
@@ -208,3 +199,30 @@ def __default__(self, data, children, meta):
         raise NotImplementedError(
             f"Calling __default__, i.e., unknown grammar concept. data: {data}, children: {children}, meta: {meta}"
         )
+
+    def _recursive_expression_phrase(self, arg):
+        """ Helper function for parsing `expression_phrase`. Recursively sorts out
+        the correct precedence for $and, $or and $nor.
+
+        """
+        if len(arg) == 1:
+            # without NOT
+            return arg[0]
+
+        # handle the case of {"$not": {"$or": [expr1, expr2]}} using {"$nor": [expr1, expr2]}.
+        if "$or" in arg[1]:
+            return {"$nor": self._recursive_expression_phrase([arg[1]["$or"]])}
+
+        # handle the case of {"$not": {"$and": [expr1, expr2]}} using per-expression negation,
+        # e.g. {"$and": [{prop1: {"$not": expr1}}, {prop2: {"$not": ~expr2}}]}.
+        # Note that this is not the same as NOT (expr1 AND expr2)!
+        if "$and" in arg[1]:
+            return {
+                "$and": [
+                    self._recursive_expression_phrase(["NOT", subdict])
+                    for subdict in arg[1]["$and"]
+                ]
+            }
+
+        # simple case of negating one expression, from NOT (expr) to ~expr.
+        return {prop: {"$not": expr} for prop, expr in arg[1].items()}
diff --git a/tests/filtertransformers/test_mongo.py b/tests/filtertransformers/test_mongo.py
@@ -226,6 +226,63 @@ def test_operators(self):
         # OPTIONAL
         # self.assertEqual(self.transform("((NOT (_exmpl_a>_exmpl_b)) AND _exmpl_x>0)"), {})
 
+        self.assertEqual(
+            self.transform("NOT (a>1 AND b>1)"),
+            {"$and": [{"a": {"$not": {"$gt": 1}}}, {"b": {"$not": {"$gt": 1}}}]},
+        )
+
+        self.assertEqual(
+            self.transform("NOT (a>1 AND b>1 OR c>1)"),
+            {
+                "$nor": [
+                    {"$and": [{"a": {"$gt": 1}}, {"b": {"$gt": 1}}]},
+                    {"c": {"$gt": 1}},
+                ]
+            },
+        )
+
+        self.assertEqual(
+            self.transform("NOT (a>1 AND ( b>1 OR c>1 ))"),
+            {
+                "$and": [
+                    {"a": {"$not": {"$gt": 1}}},
+                    {"$nor": [{"b": {"$gt": 1}}, {"c": {"$gt": 1}}]},
+                ]
+            },
+        )
+
+        self.assertEqual(
+            self.transform("NOT (a>1 AND ( b>1 OR (c>1 AND d>1 ) ))"),
+            {
+                "$and": [
+                    {"a": {"$not": {"$gt": 1}}},
+                    {
+                        "$nor": [
+                            {"b": {"$gt": 1}},
+                            {"$and": [{"c": {"$gt": 1}}, {"d": {"$gt": 1}}]},
+                        ]
+                    },
+                ]
+            },
+        )
+
+        self.assertEqual(
+            self.transform(
+                'elements HAS "Ag" AND NOT ( elements HAS "Ir" AND elements HAS "Ac" )'
+            ),
+            {
+                "$and": [
+                    {"elements": {"$in": ["Ag"]}},
+                    {
+                        "$and": [
+                            {"elements": {"$not": {"$in": ["Ir"]}}},
+                            {"elements": {"$not": {"$in": ["Ac"]}}},
+                        ]
+                    },
+                ]
+            },
+        )
+
         self.assertEqual(self.transform("5 < 7"), {7: {"$gt": 5}})
 
         with self.assertRaises(VisitError):