def time_code(timing_data_obj, code_block_name, is_main_timer=False):
"""Starts timer at entry, stores elapsed time at exit
If `is_main_timer=True`, the start time is stored in the timing_data_obj,
allowing calculation of total elapsed time 'on the fly' (e.g. to enforce
a time limit) using `get_main_elapsed_time(timing_data_obj)`.
"""
start_time = timing.default_timer()
if is_main_timer:
timing_data_obj.main_timer_start_time = start_time
yield
elapsed_time = timing.default_timer() - start_time
prev_time = timing_data_obj.get(code_block_name, 0)
timing_data_obj[code_block_name] = prev_time + elapsed_time
def get_main_elapsed_time(timing_data_obj):
"""Returns the time since entering the main `time_code` context"""
current_time = timing.default_timer()
try:
return current_time - timing_data_obj.main_timer_start_time
except AttributeError as e:
if 'main_timer_start_time' in str(e):
raise e from AttributeError(
"You need to be in a 'time_code' context to use `get_main_elapsed_time()`."
)
