nac3-spec/hm-inference/test.py

72 lines
1.2 KiB
Python

from __future__ import annotations
import ast
from ast_visitor import Visitor
from nac3_types import *
from primitives import *
src = """
a = 1 // 1
if 1 <= 1 < 2:
pass
b = test_virtual(virtual(bar, Foo))
b = test_virtual(virtual(foo, Foo))
"""
foo = TObj('Foo', {
'a': TInt,
}, [])
foo2 = TObj('Foo2', {
'a': TInt,
}, [])
bar = TObj('Bar', {
'a': TInt,
'b': TInt
}, [], [foo])
type_mapping = {
'Foo': foo,
'Foo2': foo2,
'Bar': bar,
}
prelude = {
'foo': foo,
'foo2': foo2,
'bar': bar,
'test_virtual': TFunc([FuncArg('a', TVirtual(foo), False)], TInt, [])
}
print('-----------')
print('prelude')
for key, value in prelude.items():
print(f'{key}: {value}')
print('-----------')
v = Visitor(src, prelude.copy(), lambda x: type_mapping[x])
print(src)
v.visit(ast.parse(src))
for a, b in v.virtuals:
assert isinstance(a, TObj)
assert b is a or b in a.parents
print('-----------')
print('calls')
for x in v.calls:
x.check()
print(f'{x.find()}')
print('-----------')
print('assignments')
for key, value in v.assignments.items():
if key not in prelude:
value.check()
print(f'{key}: {value.find()}')