From 8a77dcae71e95e70bb1c50241fa6ccce70ed943f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Thu, 16 May 2024 21:57:33 +0100
Subject: [PATCH] Add `and` `or` `not` Add `iif(v, t, f)` Add `==` `!=` `<`
 `<=` `>` `>=`

---
 py/math_expression.py | 43 ++++++++++++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/py/math_expression.py b/py/math_expression.py
index 08eebe9..95dfdec 100644
--- a/py/math_expression.py
+++ b/py/math_expression.py
@@ -16,7 +16,9 @@ operators = {
     ast.BitAnd: op.and_,
     ast.BitOr: op.or_,
     ast.Invert: op.invert,
-    ast.And: op.and_,
+    ast.And: lambda a, b: 1 if a and b else 0,
+    ast.Or: lambda a, b: 1 if a or b else 0,
+    ast.Not: lambda a: 0 if a else 1,
 }
 
 # TODO: restructure args to provide more info, generate hint based on args to save duplication
@@ -66,6 +68,11 @@ functions = {
         "call": lambda a = None: int(a),
         "hint": "number"
     },
+    "iif": {
+        "args": (3, 3),
+        "call": lambda a, b, c = None: b if a else c,
+        "hint": "value, truepart, falsepart"
+    },
 }
 
 autocompleteWords = list({
@@ -150,15 +157,20 @@ class MathExpression:
 
         lookup = {"a": a, "b": b, "c": c}
 
+        def eval_op(l, r):
+            l = eval_expr(l)
+            r = eval_expr(r)
+            l = l if isinstance(l, int) else float(l)
+            r = r if isinstance(r, int) else float(r)
+            return operators[type(node.op)](l, r)
+
         def eval_expr(node):
-            if isinstance(node, ast.Num):
+            if isinstance(node, ast.Constant) or isinstance(node, ast.Num):
                 return node.n
             elif isinstance(node, ast.BinOp):
-                l = eval_expr(node.left)
-                r = eval_expr(node.right)
-                l = l if isinstance(l, int) else float(l)
-                r = r if isinstance(r, int) else float(r)
-                return operators[type(node.op)](l, r)
+                return eval_op(node.left, node.right)
+            elif isinstance(node, ast.BoolOp):
+                return eval_op(node.values[0], node.values[1])
             elif isinstance(node, ast.UnaryOp):
                 return operators[type(node.op)](eval_expr(node.operand))
             elif isinstance(node, ast.Attribute):
@@ -192,6 +204,23 @@ class MathExpression:
                         args.append(eval_expr(arg))
                     return fn["call"](*args)
                 raise NameError(f"Invalid function call: {node.func.id}")
+            elif isinstance(node, ast.Compare):
+                l = eval_expr(node.left)
+                r = eval_expr(node.comparators[0])
+                if isinstance(node.ops[0], ast.Eq):
+                    return 1 if l == r else 0
+                if isinstance(node.ops[0], ast.NotEq):
+                    return 1 if l != r else 0
+                if isinstance(node.ops[0], ast.Gt):
+                    return 1 if l > r else 0
+                if isinstance(node.ops[0], ast.GtE):
+                    return 1 if l >= r else 0
+                if isinstance(node.ops[0], ast.Lt):
+                    return 1 if l < r else 0
+                if isinstance(node.ops[0], ast.LtE):
+                    return 1 if l <= r else 0
+                raise NotImplementedError(
+                    "Operator " + node.ops[0].__class__.__name__ + " not supported.")
             else:
                 raise TypeError(node)