Coverage for src/sensai/util/profiling.py: 29%

1import functools 

2import logging 

3from typing import Optional, Callable 

4from time import time 

5 

6 

7log = logging.getLogger(__name__) 

8 

9 

10def profiled(func: Optional[Callable] = None, *, sampling_interval_seconds: Optional[float] = None, 

11 print_to_console=True, log_to_file=False, file_name: Optional[str] = None): 

12 """ 

13 Function decorator for profiling the annotated/given function with the pyinstrument profiler 

14 

15 :param func: the function to be profiled 

16 :param sampling_interval_seconds: sampling every <sampling_interval_seconds> seconds, if None use default parameter of pyinstrument 

17 profiler 

18 :param print_to_console: if, true print profiler output to console 

19 :param log_to_file: if true, logs profiler output to html file 

20 :param file_name: optional file name for log, if None defaults to file name derived from function name 

21 :return: the wrapped function 

22 """ 

23 from pyinstrument import Profiler 

24 

25 def wrapper(_func): 

26 @functools.wraps(_func) 

27 def _wrapper(*args, **kwargs): 

28 _profiler = Profiler(sampling_interval_seconds) if sampling_interval_seconds is not None else Profiler() 

29 _profiler.start() 

30 result = _func(*args, **kwargs) 

31 _profiler.stop() 

32 if print_to_console: 

33 print(_profiler.output_text(unicode=True, color=True)) 

34 if log_to_file: 

35 name = file_name if file_name is not None else _func.__name__ + ".profile" 

36 with open(f"{name}.html", "w") as f: 

37 f.write(_profiler.output_html()) 

38 return result 

39 return _wrapper 

40 

41 if func: 

42 return wrapper(func) 

43 

44 return wrapper 

45 

46 

47def timed(func): 

48 """ 

49 Function decorator which logs execution times of the wrapped/annotated function at INFO level 

50 :param func: the function whose execution time to log 

51 :return: the wrapped function 

52 """ 

53 @functools.wraps(func) 

54 def _wrapper(*args, **kwargs): 

55 start = int(round(time() * 1000)) 

56 try: 

57 return func(*args, **kwargs) 

58 finally: 

59 end_ = int(round(time() * 1000)) - start 

60 log.info(f"Execution time of {func.__qualname__}: {end_ if end_ > 0 else 0} ms") 

61 return _wrapper