Support None and loose identity operators (#11)

Co-authored-by: memona <[email protected]>
Co-authored-by: Aarni Koskela <[email protected]>

leval/evaluator.py

@@ -7,14 +7,23 @@
     InvalidConstant,
     InvalidNode,
     InvalidOperation,
+    NoSuchValue,
     Timeout,
     TooComplex,
 )
 from leval.universe.base import BaseEvaluationUniverse
 from leval.utils import expand_name
 
+try:
+    from types import NoneType
+except ImportError:
+    NoneType = type(None)  # type: ignore
+
+
 DEFAULT_ALLOWED_CONTAINER_TYPES = frozenset((tuple, set))
-DEFAULT_ALLOWED_CONSTANT_TYPES = frozenset((str, int, float))
+DEFAULT_ALLOWED_CONSTANT_TYPES = frozenset(
+    (str, int, float, NoneType),
+)
 
 
 def _default_if_none(value, default):
@@ -104,12 +113,22 @@ def visit(self, node):  # noqa: D102
         finally:
             self.depth -= 1
 
+    def _visit_or_none(self, value: ast.AST) -> Any:
+        try:
+            return self.visit(value)
+        except NoSuchValue:
+            return None
+
     def visit_Compare(self, node):  # noqa: D102
         if len(node.ops) != 1:
             raise InvalidOperation("Only simple comparisons are supported", node=node)
-        left = self.visit(node.left)
-        right = self.visit(node.comparators[0])
         op = node.ops[0]
+        if isinstance(op, (ast.Is, ast.IsNot)):
+            left = self._visit_or_none(node.left)
+            right = self._visit_or_none(node.comparators[0])
+        else:
+            left = self.visit(node.left)
+            right = self.visit(node.comparators[0])
         return self.universe.evaluate_binary_op(op, left, right)
 
     def visit_Call(self, node):  # noqa: D102

leval/rewriter_evaluator.py

@@ -14,6 +14,7 @@
     "and",
     "for",
     "in",
+    "is",
     "not",
     "or",
 }

leval/universe/default.py

@@ -38,6 +38,8 @@ def binop(a, b):
     ast.LtE: operator.le,
     ast.In: lambda a, b: a in b,
     ast.NotIn: lambda a, b: a not in b,
+    ast.Is: operator.is_,
+    ast.IsNot: operator.is_not,
 }
 
 

leval/universe/weakly_typed.py

@@ -55,6 +55,8 @@ class WeaklyTypedEvaluationUniverse(EvaluationUniverse):
         ast.LtE: weakly_typed_operation(operator.le),
         ast.In: lambda a, b: a in b,
         ast.NotIn: lambda a, b: a not in b,
+        ast.Is: operator.is_,
+        ast.IsNot: operator.is_not,
     }
 
 

leval_tests/test_leval.py

@@ -40,9 +40,19 @@
     ("complex comparison", "(8, 3) >= (9, 2)", False),
     ("precedence (no parens here!)", "4 + 3 * 5 + 2", 21),
     ("precedence (with parens)", "4 + 3 * (5 + 2)", 25),
+    ("identity operator: is not", "foo is not None", True),
+    (
+        "identity operator: Is ((variable is present in values as None)",
+        "never is None",
+        True,
+    ),
+    ("identity operator: Is (variable is not present in values)", "abc is None", True),
+    ("identity operator: Is (variable is not present in values)", "abc is 4", False),
+    ("identity operator: Is (variable is present in values)", "bar is None", False),
 ]
 
 error_cases = [
+    ("Undefined variable comparison", "abc > 5", NoSuchValue),
     ("Can't index", "indx[9]", InvalidOperation),
     ("Can't call arbitrary functions", '__import__("sys")', NoSuchFunction),
     ("Invalid name", "njep + nop", NoSuchValue),