For code generators make functions assign aware (incomplete experiment)

assign_aware.py

import inspect
import ast


def get_assign_names():
    names = [] 
    caller_frame = inspect.stack()[2]
    call = ast.parse(caller_frame.code_context[0])
    assign = None
    for stmt in call.body:
        if isinstance(stmt, ast.Assign):
            assign = stmt
    if assign is None:
        return []
    for target in assign.targets:
        if isinstance(target, ast.Tuple):
            for elt in target.elts:
                if isinstance(elt, ast.Name):
                    names.append(elt.id)
        elif isinstance(target, ast.Name):
            names.append(target.id)
    return names

class AssignAware(ast.NodeTransformer):
    def visit_Assign(self, assign):
        glob = globals()
        names = []
        if isinstance(assign.value, ast.Call):
            func_name = assign.value.func.id
            args = inspect.getargspec(glob[func_name])[0]
            if args[0] != "assign":
                return assign
            for target in assign.targets:
                if isinstance(target, ast.Tuple):
                    for elt in target.elts:
                        if isinstance(elt, ast.Name):
                            names.append(elt.id)
                elif isinstance(target, ast.Name):
                    names.append(target.id)
            assign.value.args.insert(0, ast.Tuple(
                ctx=ast.Load(),
                elts=[ast.Str(s=x) for x in names]
            ))
        return assign


def assign_aware(function):
    source = inspect.getsource(function)
    ast_source = ast.parse(source)
    decorator_list = []
    for decorator in ast_source.body[0].decorator_list:
        if isinstance(decorator, ast.Name):
            if decorator.id != "assign_aware":
                decorator_list.append(decorator)
    ast_source.body[0].decorator_list = decorator_list
    ast_source = AssignAware().visit(ast_source)
    ast_source = ast.fix_missing_locations(ast_source)
    func = compile(
        ast_source,
        inspect.getfile(function),
        'exec'
    )
    env = dict(globals())
    exec(func, env)
    return env[function.__name__]


def call(assign):
    return assign

class Huhu(object):
    def bla(self, assign):
        return assign

@assign_aware
def test():
    (
        a,
        b
    ) = call()
    print(a, b)
    # (
    #     a,
    #     b
    # ) = Huhu().bla()
    # print(a, b)

test()