From 531f4e3576c71976920ce54b93abd6e3cefa2c59 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 29 Aug 2021 13:49:04 +0200 Subject: [PATCH] Add `Profile()` profiler (#4587) * Add `Profile()` profiler * CamelCase Timeout --- utils/general.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/utils/general.py b/utils/general.py index 16244903575a..c74d8bb299de 100755 --- a/utils/general.py +++ b/utils/general.py @@ -39,8 +39,17 @@ os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads -class timeout(contextlib.ContextDecorator): - # Usage: @timeout(seconds) decorator or 'with timeout(seconds):' context manager +class Profile(contextlib.ContextDecorator): + # Usage: @Profile() decorator or 'with Profile():' context manager + def __enter__(self): + self.start = time.time() + + def __exit__(self, type, value, traceback): + print(f'Profile results: {time.time() - self.start:.5f}s') + + +class Timeout(contextlib.ContextDecorator): + # Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True): self.seconds = int(seconds) self.timeout_message = timeout_msg