parser: support chained binary expression

invariantlabs-ai · Jul 5, 2024 · 0576332 · 0576332
1 parent b3f8514
commit 0576332
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 1 deletion.
diff --git a/invariant/language/parser.py b/invariant/language/parser.py
@@ -36,7 +36,7 @@
     term: factor TERM_OPERATOR factor | factor
     factor: power FACTOR_OPERATOR power | power
     power: atom POWER_OPERATOR atom | atom
-    atom: unary_expr | NUMBER | multiline_string | STRING | ID | "(" expr ")" | member_access | key_access | func_call | typed_identifier | tool_ref | object_literal | list_literal | STAR | value_ref
+    atom: unary_expr | NUMBER | multiline_string | STRING | ID | "(" expr ")" | member_access | key_access | func_call | typed_identifier | tool_ref | object_literal | list_literal | STAR | value_ref | expr
                    
     unary_expr: UNARY_OPERATOR expr
     func_call: expr  "(" ( (expr ("," expr)*)? ("," kwarg ("," kwarg)*)? ) ")" | \

diff --git a/tests/test_parser.py b/tests/test_parser.py
@@ -306,5 +306,65 @@ def test_unary_func_call_precedence(self):
         self.assertIsInstance(policy.statements[1].body[1], ast.UnaryExpr)
         self.assertIsInstance(policy.statements[1].body[1].expr, ast.FunctionCall)
 
+    def test_multi_binary(self):
+        policy = parse("""
+        raise "found result" if:
+            (output: ToolOutput)
+            a := "Cheerio" + "1" + "test"
+            print(a)
+        """)
+        self.assertIsInstance(policy.statements[0].body[1], ast.BinaryExpr)
+        # left is id
+        self.assertIsInstance(policy.statements[0].body[1].left, ast.Identifier)
+        # right is binary
+        self.assertIsInstance(policy.statements[0].body[1].right, ast.BinaryExpr)
+        rhs = policy.statements[0].body[1].right
+        # left or rhs is binary
+        self.assertIsInstance(rhs.left, ast.BinaryExpr)
+        self.assertIsInstance(rhs.right, ast.StringLiteral)
+        # left is id
+        self.assertIsInstance(rhs.left.left, ast.StringLiteral)
+        # right is id
+        self.assertIsInstance(rhs.left.right, ast.StringLiteral)
+
+    def test_in_with_member(self):
+        policy = parse("""
+        raise "found result" if:
+            (output: ToolOutput)
+            "abc" + "efg" in output.content
+        """)
+
+        self.assertIsInstance(policy.statements[0].body[1], ast.BinaryExpr)
+        self.assertIsInstance(policy.statements[0].body[1].left, ast.BinaryExpr)
+
+        self.assertIsInstance(policy.statements[0].body[1].left.left, ast.StringLiteral)
+        self.assertIsInstance(policy.statements[0].body[1].left.right, ast.StringLiteral)
+
+        self.assertIsInstance(policy.statements[0].body[1].right, ast.MemberAccess)
+        self.assertIsInstance(policy.statements[0].body[1].right.expr, ast.Identifier)
+        self.assertEqual(policy.statements[0].body[1].right.member, "content")
+        self.assertEqual(policy.statements[0].body[1].op, "in")
+
+    def test_in_with_member_three_components(self):
+        policy = parse("""
+        raise "found result" if:
+            (output: ToolOutput)
+            "abc" + "efg" + "hij" in output.content
+        """)
+
+        self.assertIsInstance(policy.statements[0].body[1], ast.BinaryExpr)
+        self.assertIsInstance(policy.statements[0].body[1].left, ast.BinaryExpr)
+
+        self.assertIsInstance(policy.statements[0].body[1].left.left, ast.BinaryExpr)
+        self.assertIsInstance(policy.statements[0].body[1].left.right, ast.StringLiteral)
+
+        self.assertIsInstance(policy.statements[0].body[1].left.left.left, ast.StringLiteral)
+        self.assertIsInstance(policy.statements[0].body[1].left.left.right, ast.StringLiteral)
+
+        self.assertIsInstance(policy.statements[0].body[1].right, ast.MemberAccess)
+        self.assertIsInstance(policy.statements[0].body[1].right.expr, ast.Identifier)
+        self.assertEqual(policy.statements[0].body[1].right.member, "content")
+
+
 if __name__ == "__main__":
     unittest.main()