baseimage-tutorial-framework/lib/source_code.py

import ast
import inspect
import re
from io import StringIO


def find_local_variable_value(func, local_variable_name):
    func_src = inspect.getsource(func)
    func_ast = ast.parse(func_src)
    for node in ast.walk(func_ast):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name):
                    if target.id == local_variable_name:
                        return node.value.s


def get_source_code(func, strip_comments=False):
    source = inspect.getsource(func)
    if strip_comments:
        # TODO: less fragile way to do this (tokenizer sadly inserts whitespaces all around)
        comment_pattern = re.compile('^(\s.*)#.*$')
        source = ''.join(
            line for line in StringIO(source).readlines() if re.match(comment_pattern, line) is None
        )
    return source