custom decorator for writing profiling information to a file

tests.py

def profile_memory(output):
    if isinstance(output, str):
        output = open(output, 'w')

    def decorator(func):
        @profile(stream=output)
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator