fixed #18 : I can evaluate concept

tests/common/test_ast_utils.py

@@ -2,7 +2,8 @@ import ast
 
 import pytest
 
-from common.ast_utils import NamesWithAttributesVisitor, UnreferencedNamesVisitor, UnreferencedVariablesVisitor
+from common.ast_utils import NamesWithAttributesVisitor, UnreferencedNamesVisitor, UnreferencedVariablesVisitor, \
+    WhereConstraintVisitor
 
 
 @pytest.mark.parametrize("source, expected", [
@@ -15,9 +16,9 @@ from common.ast_utils import NamesWithAttributesVisitor, UnreferencedNamesVisito
     ("func(x=a, y=b)", {"func", "a", "b"}),
 
 ])
-def test_i_can_get_unreferenced_names_from_simple_expressions(context, source, expected):
+def test_i_can_get_unreferenced_names_from_simple_expressions(source, expected):
     ast_ = ast.parse(source)
-    visitor = UnreferencedNamesVisitor(context)
+    visitor = UnreferencedNamesVisitor()
     visitor.visit(ast_)
 
     assert visitor.names == expected
@@ -42,9 +43,58 @@ def test_name_with_attribute():
     ("for i in range(10): pass", set()),
     ("func(x=a, y=b)", {"a", "b", "x", "y"}),
 ])
-def test_i_can_get_unreferenced_variables_from_simple_expressions(context, source, expected):
+def test_i_can_get_unreferenced_variables_from_simple_expressions(source, expected):
     ast_ = ast.parse(source)
-    visitor = UnreferencedVariablesVisitor(context)
+    visitor = UnreferencedVariablesVisitor()
     visitor.visit(ast_)
 
     assert visitor.names == expected
+
+
+def test_i_can_get_where_constraints():
+    expr = "a == 2"
+    ast_tree = ast.parse(expr, f"<user input>", 'eval')
+    visitor = WhereConstraintVisitor(ast_tree)
+
+    res = visitor.get_constraints()
+    assert res == {"a": [WhereConstraintVisitor.WhereConstraint("a == 2")]}
+    assert isinstance(res["a"][0].ast_tree, ast.Expression)
+
+
+def test_i_can_get_where_constraints_when_and():
+    expr = "a == 2 and isinstance(b, Concept)"
+    ast_tree = ast.parse(expr, f"<user input>", 'eval')
+    visitor = WhereConstraintVisitor(ast_tree)
+
+    res = visitor.get_constraints()
+    assert res == {"a": [WhereConstraintVisitor.WhereConstraint("a == 2")],
+                   "b": [WhereConstraintVisitor.WhereConstraint("isinstance(b, Concept)")],
+                   "Concept": [WhereConstraintVisitor.WhereConstraint("isinstance(b, Concept)")]}
+
+    assert isinstance(res["a"][0].ast_tree, ast.Expression)
+    assert isinstance(res["b"][0].ast_tree, ast.Expression)
+    assert isinstance(res["Concept"][0].ast_tree, ast.Expression)
+
+
+def test_i_can_get_where_constraint_when_name_is_reference_several_times():
+    expr = "isinstance(a, int) and a == 2"
+    ast_tree = ast.parse(expr, f"<user input>", 'eval')
+    visitor = WhereConstraintVisitor(ast_tree)
+
+    res = visitor.get_constraints()
+    assert res == {"a": [WhereConstraintVisitor.WhereConstraint("isinstance(a, int)"),
+                         WhereConstraintVisitor.WhereConstraint("a == 2")],
+                   "int": [WhereConstraintVisitor.WhereConstraint("isinstance(a, int)")]}
+
+    assert isinstance(res["a"][0].ast_tree, ast.Expression)
+    assert isinstance(res["a"][1].ast_tree, ast.Expression)
+    assert isinstance(res["int"][0].ast_tree, ast.Expression)
+
+
+def test_i_cannot_get_constraint_when_or():
+    expr = "isinstance(a, int) or a == 2"
+    ast_tree = ast.parse(expr, f"<user input>", 'eval')
+    visitor = WhereConstraintVisitor(ast_tree)
+
+    with pytest.raises(NotImplementedError):
+        visitor.get_constraints()