Added a Lark to Django Query converter

optimade/filtertransformers/django.py

@@ -0,0 +1,103 @@
+import operator
+from optimade.filterparser import LarkParser
+from lark import Tree
+from lark.lexer import Token
+from django.db.models import Q
+import json
+import os
+
+
+class DjangoQueryError(Exception):
+    pass
+
+
+django_db_keys = {
+    "natoms": "entry__natoms",
+    "chemical_formula": "composition__formula__in",
+    "stability": "stability",
+    "band_gap": "calculation__band_gap",
+    "element": "composition__element_list__contains",
+    "nelements": "entry__composition__ntypes",
+}
+
+
+class Lark2Django:
+    def __init__(self):
+        self.opers = {
+            "=": self.eq,
+            ">": self.gt,
+            ">=": self.ge,
+            "<": self.lt,
+            "<=": self.le,
+            "!=": self.ne,
+            "OR": self.or_,
+            "AND": self.and_,
+            "NOT": self.not_,
+        }
+        self.parser = LarkParser(version=(0, 9, 7))
+
+    def parse_raw_q(self, raw_query):
+        return self.parser.parse(raw_query)
+
+    def eq(self, a, b):
+        return Q(**{a: b})
+
+    def gt(self, a, b):
+        return Q(**{a + "__gt": b})
+
+    def ge(self, a, b):
+        return Q(**{a + "__gte": b})
+
+    def lt(self, a, b):
+        return Q(**{a + "__lt": b})
+
+    def le(self, a, b):
+        return Q(**{a + "__lte": b})
+
+    def ne(self, a, b):
+        return ~Q(**{a: b})
+
+    def not_(self, a):
+        return ~a
+
+    def and_(self, a, b):
+        return operator.and_(a, b)
+
+    def or_(self, a, b):
+        return operator.or_(a, b)
+
+    def evaluate(self, parse_Tree):
+        if isinstance(parse_Tree, Tree):
+            children = parse_Tree.children
+            if len(children) == 1:
+                return self.evaluate(children[0])
+            elif len(children) == 2:
+                op_fn = self.evaluate(children[0])
+                return op_fn(self.evaluate(children[1]))
+            elif len(children) == 3:
+                if parse_Tree.data == "comparison":
+                    db_prop = self.evaluate(children[0])
+                    op_fn = self.evaluate(children[1])
+
+                    if db_prop in django_db_keys.keys():
+                        return op_fn(
+                            django_db_keys[db_prop], self.evaluate(children[2])
+                        )
+                    else:
+                        raise DjangoQueryError(
+                            "Unknown property is queried : " + (db_prop)
+                        )
+
+                else:
+                    op_fn = self.evaluate(children[1])
+                    return op_fn(self.evaluate(children[0]), self.evaluate(children[2]))
+            else:
+                raise DjangoQueryError("Not compatible format. Tree has >3 children")
+
+        elif isinstance(parse_Tree, Token):
+            if parse_Tree.type == "VALUE":
+                return parse_Tree.value
+            elif parse_Tree.type in ["NOT", "CONJUNCTION", "OPERATOR"]:
+                return self.opers[parse_Tree.value]
+        else:
+            raise DjangoQueryError("Not a Lark Tree or Token")

optimade/filtertransformers/tests/test_django.py

@@ -0,0 +1,28 @@
+import json
+from optimade.filtertransformers.django import Lark2Django
+from unittest import TestCase
+import os
+
+test_data = [
+    ("band_gap<1", "(AND: ('calculation__band_gap__lt', '1'))"),
+    (
+        "natoms >= 8 OR nelements<5 AND stability>=0.5",
+        "(AND: (OR: ('entry__natoms__gte', '8'), ('entry__composition__ntypes__lt', '5')), ('stability__gte', '0.5'))",
+    ),
+    ("NOT natoms >= 8", "(NOT (AND: ('entry__natoms__gte', '8')))"),
+    (
+        "element = Th AND element != Na ",
+        "(AND: ('composition__element_list__contains', 'Th'), (NOT (AND: ('composition__element_list__contains', 'Na'))))",
+    ),
+]
+
+
+class DjangoParserTest(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.Transformer = Lark2Django()
+
+    def test_query_conversion(self):
+        for raw_q, dj_q in test_data:
+            parsed_tree = self.Transformer.parse_raw_q(raw_q)
+            self.assertEqual(str(self.Transformer.evaluate(parsed_tree)), dj_q)

requirements.txt

@@ -4,3 +4,4 @@ lark-parser==0.5.6
 mongogrant==0.2.2
 mongomock==3.16.0
 pymongo==3.8.0
+django==2.2.5