Attach function/module source on import via a decorator

source_recorder.py

"""Attach function/module source on import via a decorator
"""

import sys
import inspect


SOURCE = '_source'


def get_source(f):
    return getattr(f, SOURCE)


def set_source(f, source):
    setattr(f, SOURCE, source)


def has_source(f):
    return hasattr(f, SOURCE)


def _attach_source(f):
    source = inspect.getsource(f)
    set_source(f, source)


def attach_source(f, record_module=True):
    if record_module:
        module = sys.modules[f.__module__]
        if not has_source(module):
            _attach_source(module)
    _attach_source(f)
    return f