101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
import ast
|
|
|
|
import pytest
|
|
|
|
from common.ast_utils import NamesWithAttributesVisitor, UnreferencedNamesVisitor, UnreferencedVariablesVisitor, \
|
|
WhereConstraintVisitor
|
|
|
|
|
|
@pytest.mark.parametrize("source, expected", [
|
|
("a,b", {"a", "b"}),
|
|
("isinstance(a, int)", {"isinstance", "a", "int"}),
|
|
("date.today()", {"date"}),
|
|
("test()", {"test"}),
|
|
("sheerka.test()", {"sheerka"}),
|
|
("for i in range(10): pass", set()),
|
|
("func(x=a, y=b)", {"func", "a", "b"}),
|
|
|
|
])
|
|
def test_i_can_get_unreferenced_names_from_simple_expressions(source, expected):
|
|
ast_ = ast.parse(source)
|
|
visitor = UnreferencedNamesVisitor()
|
|
visitor.visit(ast_)
|
|
|
|
assert visitor.names == expected
|
|
|
|
|
|
def test_name_with_attribute():
|
|
# Looks for all attributes for a given name
|
|
ast_ = ast.parse("foo.bar.baz", "<src>", mode="exec")
|
|
assert NamesWithAttributesVisitor().get_sequences(ast_, "foo") == [["foo", "bar", "baz"]]
|
|
|
|
# It parses all expressions / statements
|
|
ast_ = ast.parse("foo.bar.baz; one.two.three; foo.bar", "<src>", mode="exec")
|
|
assert NamesWithAttributesVisitor().get_sequences(ast_, "foo") == [["foo", "bar", "baz"], ["foo", "bar"]]
|
|
|
|
|
|
@pytest.mark.parametrize("source, expected", [
|
|
("a,b", {"a", "b"}),
|
|
("isinstance(a, int)", {"a", "int"}),
|
|
("date.today()", set()),
|
|
("test()", set()),
|
|
("sheerka.test()", set()),
|
|
("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(source, expected):
|
|
ast_ = ast.parse(source)
|
|
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()
|