05360171创建于 2022年3月18日历史提交
# Copyright 2021 Huawei Technologies Co., Ltd

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

import time

from collections import defaultdict



_total_times = defaultdict(lambda:  0)

_start_times = defaultdict(lambda: -1)

_disabled_names = set()

_timer_stack = []

_running_timer = None

_disable_all = False



def disable_all():

	global _disable_all

	_disable_all = True



def enable_all():

	global _disable_all

	_disable_all = False



def disable(fn_name):

	""" Disables the given function name fom being considered for the average or outputted in print_stats. """

	_disabled_names.add(fn_name)



def enable(fn_name):

	""" Enables function names disabled by disable. """

	_disabled_names.remove(fn_name)



def reset():

	""" Resets the current timer. Call this at the start of an iteration. """

	global _running_timer

	_total_times.clear()

	_start_times.clear()

	_timer_stack.clear()

	_running_timer = None



def start(fn_name, use_stack=True):

	"""

	Start timing the specific function.

	Note: If use_stack is True, only one timer can be active at a time.

	      Once you stop this timer, the previous one will start again.

	"""

	global _running_timer, _disable_all

	

	if _disable_all:

		return



	if use_stack:

		if _running_timer is not None:

			stop(_running_timer, use_stack=False)

			_timer_stack.append(_running_timer)

		start(fn_name, use_stack=False)

		_running_timer = fn_name

	else:

		_start_times[fn_name] = time.perf_counter()



def stop(fn_name=None, use_stack=True):

	"""

	If use_stack is True, this will stop the currently running timer and restore

	the previous timer on the stack if that exists. Note if use_stack is True,

	fn_name will be ignored.



	If use_stack is False, this will just stop timing the timer fn_name.

	"""

	global _running_timer, _disable_all



	if _disable_all:

		return



	if use_stack:

		if _running_timer is not None:

			stop(_running_timer, use_stack=False)

			if len(_timer_stack) > 0:

				_running_timer = _timer_stack.pop()

				start(_running_timer, use_stack=False)

			else:

				_running_timer = None

		else:

			print('Warning: timer stopped with no timer running!')

	else:

		if _start_times[fn_name] > -1:

			_total_times[fn_name] += time.perf_counter() - _start_times[fn_name]

		else:

			print('Warning: timer for %s stopped before starting!' % fn_name)





def print_stats():

	""" Prints the current timing information into a table. """

	print()



	all_fn_names = [k for k in _total_times.keys() if k not in _disabled_names]



	max_name_width = max([len(k) for k in all_fn_names] + [4])

	if max_name_width % 2 == 1: max_name_width += 1

	format_str = ' {:>%d} | {:>10.4f} ' % max_name_width



	header = (' {:^%d} | {:^10} ' % max_name_width).format('Name', 'Time (ms)')

	print(header)



	sep_idx = header.find('|')

	sep_text = ('-' * sep_idx) + '+' + '-' * (len(header)-sep_idx-1)

	print(sep_text)



	for name in all_fn_names:

		print(format_str.format(name, _total_times[name]*1000))

	

	print(sep_text)

	print(format_str.format('Total', total_time()*1000))

	print()



def total_time():

	""" Returns the total amount accumulated across all functions in seconds. """ 

	return sum([elapsed_time for name, elapsed_time in _total_times.items() if name not in _disabled_names])





class env():

	"""

	A class that lets you go:

		with timer.env(fn_name):

			# (...)

	That automatically manages a timer start and stop for you.

	"""



	def __init__(self, fn_name, use_stack=True):

		self.fn_name = fn_name

		self.use_stack = use_stack



	def __enter__(self):

		start(self.fn_name, use_stack=self.use_stack)



	def __exit__(self, e, ev, t):

		stop(self.fn_name, use_stack=self.use_stack)