Sheerka-Old/tests/test_ast.py

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