simple_check.py

import inspect
import z3

def parse_type(t):
    parts = t.split('|')
    name, base = parts[0].split(':')
    return (name.strip(), base.strip(), parts[1].strip())

def constraints(func):
    sig = inspect.signature(func)
    ret = parse_type(sig.return_annotation)
    inputs = []
    var_dict = {}
    for name, param in sig.parameters.items():
        var = z3.Int(name)
        var_dict[name] = var
        constr = z3.BoolVal(True)
        annot = param.annotation
        if annot != int:
            parsed = parse_type(annot)
            var_dict[parsed[0]] = var
            constr = eval(parsed[2], locals=var_dict)
        inputs.append((name, var, constr))
    ret_var, ret_base, ret_pred = parse_type(sig.return_annotation)
    ret = z3.Int(ret_var)
    var_dict[ret_var] = ret
    ret_constr = eval(ret_pred, locals=var_dict)
    return inputs, (ret_var, ret, ret_constr)

def vc(func):
    inputs, (_, ret, constr) = constraints(func)
    premise = z3.And([c for (_, _, c) in inputs])
    args = {name:var for (name, var, _) in inputs}
    result = func(**args)
    return z3.Implies(premise, z3.substitute(constr, (ret, result)))

def check(func):
    solver = z3.Solver()
    frm = z3.simplify(z3.Not(vc(func)))
    solver.add(frm)
    if solver.check() == z3.unsat:
        return func
    raise Exception('WRONG')

@check
def foo(n: int, m: 'p : int | p > 0') -> 'z : int | z > n':
    return m + n