estimate_return_type.py

import ast
import inspect
from typing import Callable, Union, Tuple, List, Iterable, Dict, Any

METHOD_WRAPPER = type(1..__init__)
WRAPPER_DESCRIPTOR = type(int.__init__)


def l():
    return [x for x in [1, 2, 3]]


def m(x: List[int]):
    return {y(): k for k in x}


def n():
    return {"x": "y", "z": 1}


def o():
    y = [1, 2, 3, 4]
    return ["a"*(x+1) for x in y]


def p():
    return b""


def q():  # float
    x = 10 / 1.0
    return x


def r():  # str
    return "abc"[1]


def s():  # List[Union[str, int]]
    return ["10", 15, "abc", 1, 2, 3]


def t():  # List[int]
    return [10, 20]


def u():  # Tuple[int, bool]
    return 10, True


def v():  # float
    return 10.0


def w():  # str
    return "abc"


def y() -> int:  # int
    z = 1
    return z


def x():  # int
    return y()


def z():  # None
    pass


def help(obj):
    print("==========")
    print(obj)
    for k in dir(obj):
        if k.startswith("__"):
            continue
        print(k, "->", getattr(obj, k))


class Visitor:
    def __init__(self, nodes):
        self.nodes = nodes
        self.args = {}
        self.in_comp = False

    def set_function(self, f):
        self.args = f.__annotations__

    def visit(self, node):
        f = "visit_" + node.__class__.__name__
        # print(f)
        fun = getattr(self, f, self.generic_visit)
        return fun(node)

    def unique(self,  items: Iterable):
        res = []
        for item in items:
            if item not in res:
                res.append(item)

        return res

    def visit_Dict(self, node):
        keys = self.unique(map(self.visit, node.keys))
        values = self.unique(map(self.visit, node.values))
        if len(keys) > 1:
            keys = Union[tuple(keys)]
        else:
            keys = keys[0]

        if len(values) > 1:
            values = Union[tuple(values)]
        else:
            values = values[0]

        return Dict[keys, values]

    def visit_comprehension(self, node):
        types = self.visit(node.iter).__args__
        if len(types) > 1:
            return Union[types]
        else:
            return types[0]

    def visit_DictComp(self, node):
        self.in_comp = True
        self.nodes.append(node)
        # node.key and node.value are the key/value code parts
        # node.generators[0] is the iter
        key = self.visit(node.key)
        value = self.visit(node.value)
        self.in_comp = False
        return Dict[key, value]

    def visit_ListComp(self, node):
        self.in_comp = True
        self.nodes.append(node)
        # node.elt is the action happening in the listcomp
        # node.generators[0] is the iter
        target = self.visit(node.elt)
        self.in_comp = False
        return List[target]

    def visit_Bytes(self, _):
        return bytes

    def visit_Subscript(self, node):
        type_ = self.visit(node.value)
        if isinstance(type_.__getitem__, (METHOD_WRAPPER, WRAPPER_DESCRIPTOR)):
            if hasattr(type_, "__args__"):
                return type_.__args__
            return type_

        return type.__getitem__.__annotations__["return"]

    def visit_BinOp(self, node):
        cls = type(node.op)
        fname = {
            ast.Mult: "__mul__",
            ast.Div: "__truediv__",
            ast.Add: "__add__",
            ast.Sub: "__sub__"
        }[cls]
        rfname = {
            ast.Mult: "__rmul__",
            ast.Div: "__rdiv__",
            ast.Add: "__radd__",
            ast.Sub: "__rsub__"
        }[cls]

        left = self.visit(node.left)
        right = self.visit(node.right)

        real_func = getattr(left, fname, getattr(right, rfname, None))

        if left is None or left.__module__ == "typing":
            return right

        if isinstance(real_func, (METHOD_WRAPPER, WRAPPER_DESCRIPTOR)):
            # builtins/C, so make some guesses
            if issubclass(left, str):
                return str
            if issubclass(left, list):
                return left
            if issubclass(right, float):
                return float
            return left

        if real_func is not None:
            return real_func.__annotations__["return"]

    def visit_List(self, node):
        items = [self.visit(it) for it in node.elts]
        res = self.unique(items)
        if len(res) > 1:
            return List[Union[tuple(res)]]
        else:
            return List[res[0]]

    def visit_NameConstant(self, node):
        return type(node.value)

    def visit_Str(self, _):
        return str

    def visit_Tuple(self, node):
        nodes = []
        for nnode in node.elts:
            if isinstance(nnode, ast.Name):
                nodes.append(nnode.id)
            else:
                nodes.append(self.visit(nnode))

        return Tuple[tuple(nodes)]

    def visit_Name(self, node):
        var_name = node.id
        for nnode in self.nodes:
            if self.in_comp:
                if isinstance(nnode, ast.ListComp):
                    e = nnode.elt
                    if isinstance(e, ast.Name) and e == node:
                        return self.visit(nnode.generators[0])
                    # Do this too if this is a sub-sub-sub-sub element in the list
                if isinstance(nnode, ast.DictComp):
                    if nnode.key == node or nnode.value == node:
                        return self.visit(nnode.generators[0])
            elif isinstance(nnode, ast.Assign):
                targets = nnode.targets[0]
                if isinstance(targets, ast.Name):
                    if targets.id == var_name:
                        return self.visit(nnode.value)

                targets = self.visit(targets)
                value = self.visit(nnode.value)
                if not isinstance(targets, list):
                    targets = [targets]
                if not isinstance(value, list):
                    value = [value]
                for name, value in zip(targets, value):
                    if name == var_name:
                        return value
        return self.args.get(var_name, Any)

    def visit_Num(self, node):
        var_name = node._fields[0]
        return type(getattr(node, var_name))

    def visit_Call(self, node):
        f = node.func
        func: Callable = eval(f.id)
        # TODO: functions in call scope
        return func.__annotations__['return']

    def generic_visit(self, node):
        raise Exception(node.__class__.__name__)


def get_return(f):
    source = inspect.getsource(f)
    module = ast.parse(source)
    function = module.body[0]
    nodes = function.body[::-1]
    return_node = nodes[0]
    vis = Visitor(nodes)
    vis.set_function(f)
    if not isinstance(return_node, ast.Return):
        return None

    ret_val = return_node.value
    return vis.visit(ret_val)


for fun in (l, m, n, o, p, q, r, s, t, u, v, w, x, y, z):
    print(fun.__name__, "->", get_return(fun))