Coverage for postrfp / shared / utils.py: 100%

1import io 

2import time 

3import pstats 

4import logging 

5import cProfile 

6import contextlib 

7from datetime import datetime, timezone 

8from enum import Enum 

9from decimal import Decimal 

10 

11from pydantic import AnyHttpUrl 

12 

13 

14bench_log = logging.getLogger("postrfp.benchmark") 

15log = logging.getLogger(__name__) 

16 

17 

18class benchmark(object): # pragma: no cover 

19 """ 

20 Context manager for measuring and logging execution time of code blocks. 

21 

22 This utility logs the execution time with varying log levels based on how long 

23 the operation takes. Longer execution times result in higher severity log levels. 

24 

25 Parameters: 

26 name (str): Descriptive name for the operation being timed, included in log output 

27 sensitivity (float, optional): Multiplier that adjusts timing thresholds. 

28 Higher values make the benchmark more sensitive to timing differences, 

29 triggering higher log levels at shorter durations. Defaults to 1. 

30 

31 Log level thresholds: 

32 - DEBUG: time < (0.1 / sensitivity) seconds 

33 - INFO: time < (0.5 / sensitivity) seconds 

34 - WARN: time < (2.0 / sensitivity) seconds 

35 - ERROR: time ≥ (2.0 / sensitivity) seconds 

36 

37 Example: 

38 >>> with benchmark("database query", sensitivity=2): 

39 ... results = db.execute_complex_query() 

40 

41 # For a query taking 0.8 seconds with sensitivity=2: 

42 # Will log at WARN level: "database query : 0.800 seconds" 

43 """ 

44 

45 def __init__(self, name, sensitivity=1): 

46 self.name = name 

47 self.sensitivity = sensitivity 

48 

49 def __enter__(self): 

50 self.start = time.time() 

51 

52 def __exit__(self, _ty, val, _tb): 

53 elapsed = time.time() - self.start 

54 if elapsed < (0.1 / self.sensitivity): 

55 level = logging.DEBUG 

56 elif elapsed < (0.5 / self.sensitivity): 

57 level = logging.INFO 

58 elif elapsed < (2 / self.sensitivity): 

59 level = logging.WARN 

60 else: 

61 level = logging.ERROR 

62 bench_log.log(level, "%s : %0.3f seconds", self.name, elapsed) 

63 return False 

64 

65 

66@contextlib.contextmanager 

67def profiled(callers: bool = False): 

68 """ 

69 Context manager for profiling code blocks using cProfile. 

70 Usage: 

71 with profiled() as pr: 

72 # code to profile 

73 # profiling results are printed to stdout 

74 """ 

75 pr = cProfile.Profile() 

76 pr.enable() 

77 yield pr 

78 pr.disable() 

79 s = io.StringIO() 

80 ps = pstats.Stats(pr, stream=s).sort_stats("cumulative") 

81 # ps.print_stats(1000, "postrfp") 

82 ps.print_stats() 

83 if callers: 

84 ps.print_callers() 

85 print(s.getvalue()) 

86 

87 

88def json_default(obj): 

89 if isinstance(obj, set): 

90 return list(obj) 

91 elif isinstance(obj, Decimal): 

92 return float(obj) 

93 elif hasattr(obj, "as_dict"): 

94 return obj.as_dict() 

95 elif isinstance(obj, Enum): 

96 return obj.value 

97 elif isinstance(obj, AnyHttpUrl): 

98 return str(obj) 

99 else: 

100 raise TypeError("Cannot serialise %s to json." % type(obj)) 

101 

102 

103def utcnow() -> datetime: 

104 return datetime.now(timezone.utc)