Last active
September 3, 2018 09:05
-
-
Save sergio-bershadsky/3fbdc2bda456260bbfb1eb77c9371f80 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python3 | |
| import pstats | |
| import cProfile | |
| import functools | |
| class ProfileIt(object): | |
| def __call__(self, f): | |
| @functools.wraps(f) | |
| def wrapped(_self, *args, **kwargs): | |
| profiler = cProfile.Profile() | |
| profiler.enable() | |
| exception = None | |
| result = None | |
| try: | |
| result = f(_self, *args, **kwargs) | |
| except Exception as e: | |
| exception = e | |
| profiler.disable() | |
| getattr(f, "profiler_callback", self.default_callback)(_self, profiler) | |
| if exception: | |
| raise exception | |
| return result | |
| def callback(profiler_callback): | |
| f.profiler_callback = profiler_callback | |
| wrapped.callback = callback | |
| return wrapped | |
| @staticmethod | |
| def default_callback(self, profiler): | |
| s = io.StringIO() | |
| ps = pstats.Stats(profiler, stream=s).sort_stats("cumulative") | |
| ps.print_stats() | |
| print(s.getvalue()) | |
| profileit = ProfileIt() | |
| def profiler2dict(profiler): | |
| stats = pstats.Stats(profiler) | |
| return { | |
| "total_tt": stats.total_tt, | |
| "total_calls": stats.total_calls, | |
| "prim_calls": stats.prim_calls, | |
| } | |
| # | |
| # test_case.py | |
| # | |
| class ExampleTest(TestCase): | |
| @profileit | |
| def test_with_profile_and_callback(self): | |
| spass | |
| @test_with_profile_and_callback.callback | |
| def test_with_profile_callback(self, profiler): | |
| """ | |
| Here you can run asserts over profiler results | |
| """ | |
| total_calls = utils.profiler2dict(profiler)["total_calls"] | |
| self.assertGreater(total_calls, 0) | |
| @profileit | |
| def test_with_profile_print_stats_only(self): | |
| spas |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment