test_ref.py 文件源码

def profiler(request):
    class CustomProfiler(object):
        def __init__(self, request):
            if request.config.getoption("--profile"):
                import cProfile
                self.profile = cProfile.Profile()
            else:
                self.profile = None

        def __enter__(self):
            if self.profile is None:
                return False
            self.profile.enable()

        def __exit__(self, *args, **kwargs):
            if self.profile is None:
                return False
            self.profile.disable()
            self.profile.dump_stats(request.function.__name__)
            return False

    return CustomProfiler(request)