import ast import pytest from core.ast.nodes import NodeParent, GenericNodeConcept import core.ast.nodes from core.ast.visitors import ConceptNodeVisitor, UnreferencedNamesVisitor from core.builtin_concepts import BuiltinConcepts from core.sheerka import Sheerka def get_sheerka(): sheerka = Sheerka(skip_builtins_in_db=True) sheerka.initialize("mem://") return sheerka class TestNameVisitor(ConceptNodeVisitor): """ Test class for a basic Visitor test """ def __init__(self): self.names = [] def visit_Name(self, node): self.names.append(node) def test_i_can_transform_simple_ast_using_generic_node(): source = """ def my_function(a,b): for i in range(b): a = a+b return a """ tree = ast.parse(source) tree_as_concept = core.ast.nodes.python_to_concept(tree) sheerka = get_sheerka() assert tree_as_concept.node_type == "Module" assert sheerka.isinstance(tree_as_concept.get_prop("body"), BuiltinConcepts.LIST) def_func = tree_as_concept.get_prop("body")[0] assert sheerka.isinstance(def_func, BuiltinConcepts.GENERIC_NODE) assert def_func.node_type == "FunctionDef" assert def_func.parent == NodeParent(tree_as_concept, "body") assert def_func.get_prop("name") == "my_function" def_func_args = def_func.get_prop("args") assert sheerka.isinstance(def_func_args, BuiltinConcepts.GENERIC_NODE) assert def_func_args.node_type == "arguments" def_func_args_real_args = def_func_args.get_prop("args") assert sheerka.isinstance(def_func_args_real_args, BuiltinConcepts.LIST) assert len(def_func_args_real_args) == 2 assert sheerka.isinstance(def_func_args_real_args[0], BuiltinConcepts.GENERIC_NODE) assert def_func_args_real_args[0].node_type == "arg" assert def_func_args_real_args[0].parent == NodeParent(def_func_args, "args") assert def_func_args_real_args[0].get_prop("arg") == "a" assert sheerka.isinstance(def_func_args_real_args[1], BuiltinConcepts.GENERIC_NODE) assert def_func_args_real_args[1].node_type == "arg" assert def_func_args_real_args[1].parent == NodeParent(def_func_args, "args") assert def_func_args_real_args[1].get_prop("arg") == "b" def_fun_body = def_func.get_prop("body") assert sheerka.isinstance(def_fun_body, BuiltinConcepts.LIST) assert len(def_fun_body) == 2 def_fun_body_for = def_fun_body[0] assert sheerka.isinstance(def_fun_body_for, BuiltinConcepts.GENERIC_NODE) assert def_fun_body_for.node_type == "For" assert def_fun_body_for.parent == NodeParent(def_func, "body") def_fun_body_return = def_fun_body[1] assert sheerka.isinstance(def_fun_body_return, BuiltinConcepts.GENERIC_NODE) assert def_fun_body_return.node_type == "Return" assert def_fun_body_return.parent == NodeParent(def_func, "body") def test_i_can_visit_concept_node(): source = """ def my_function(a,b): for i in range(b): a = a+b return a """ node = ast.parse(source) concept_node = core.ast.nodes.python_to_concept(node) visitor = TestNameVisitor() visitor.visit(concept_node) sheerka = get_sheerka() assert sheerka.value(visitor.names[0]) == "i" assert sheerka.value(visitor.names[1]) == "range" assert sheerka.value(visitor.names[2]) == "b" assert sheerka.value(visitor.names[3]) == "a" assert sheerka.value(visitor.names[4]) == "a" assert sheerka.value(visitor.names[5]) == "b" assert sheerka.value(visitor.names[6]) == "a" def test_i_can_get_unreferenced_variables(): source = """ def my_function(a,b): for i in range(b): a = a+b return a my_function(x,y) """ sheerka = get_sheerka() node = ast.parse(source) concept_node = core.ast.nodes.python_to_concept(node) visitor = UnreferencedNamesVisitor(sheerka) visitor.visit(concept_node) values = visitor.names assert len(visitor.names) == 2 assert "x" in values assert "y" in values @pytest.mark.parametrize("source, expected", [ ("a,b", ["a", "b"]), ("isinstance(a, int)", ["a", "int"]) ]) def test_i_can_get_unreferenced_variables_from_simple_expressions(source, expected): sheerka = get_sheerka() node = ast.parse(source) concept_node = core.ast.nodes.python_to_concept(node) visitor = UnreferencedNamesVisitor(sheerka) visitor.visit(concept_node) assert sorted(list(visitor.names)) == expected def test_i_can_compare_NodeParent_with_tuple(): node_parent = NodeParent(GenericNodeConcept("For", None), "target") assert node_parent == ("For", "target") def test_i_can_transform_back(): source = """ def my_function(a,b): for i in range(b): a = a + b return a my_function(x, y) """ node = ast.parse(source) concept_node = core.ast.nodes.python_to_concept(node) transformed_back = core.ast.nodes.concept_to_python(concept_node) assert dump_ast(transformed_back) == dump_ast(node) def dump_ast(node): dump = ast.dump(node) for to_remove in [", ctx=Load()", ", kind=None", ", type_ignores=[]"]: dump = dump.replace(to_remove, "") return dump