Coverage for optimates/utils.py: 59%

35 statements  

« prev     ^ index     » next       coverage.py v7.6.7, created at 2024-11-16 11:19 -0500

1"""Utility module.""" 

2 

3import heapq 

4import logging 

5from typing import Any, Optional, TypeVar, cast 

6 

7 

8T = TypeVar('T') 

9 

10 

11########### 

12# LOGGING # 

13########### 

14 

15class Logger(logging.Logger): 

16 """Custom subclass of logging.Logger.""" 

17 

18 def __init__(self, *args: Any, **kwargs: Any) -> None: 

19 super().__init__(*args, **kwargs) 

20 self.verbosity = 0 

21 

22 def level_for_verbosity(self, verbosity: int) -> int: 

23 """Converts a verbosity level to a logging level.""" 

24 return logging.INFO - verbosity 

25 

26 def set_verbosity(self, verbosity: int) -> None: 

27 """Sets the verbosity level of the logger. 

28 A level of 0 is the normal (INFO) logging level. 

29 A higher number means more verbose.""" 

30 self.setLevel(self.level_for_verbosity(verbosity)) 

31 

32 def verbose(self, msg: str, level: int = 1) -> None: 

33 """Logs a message at the given verbosity level.""" 

34 self.log(self.level_for_verbosity(level), msg) 

35 

36 

37LOG_FMT = '%(name)s - %(message)s' 

38logging.basicConfig(format=LOG_FMT, level=logging.INFO) 

39logging.setLoggerClass(Logger) 

40logger = cast(Logger, logging.getLogger('optimates')) 

41 

42 

43######## 

44# HEAP # 

45######## 

46 

47class TopNHeap(list[T]): 

48 """Maintains the largest N elements on a heap.""" 

49 

50 def __init__(self, N: Optional[int] = None) -> None: 

51 super().__init__() 

52 self.N = N 

53 

54 def empty(self) -> bool: 

55 """Returns True if the heap is empty.""" 

56 return (len(self) == 0) 

57 

58 def top(self) -> T: 

59 """Returns the top (minimum) element on the heap.""" 

60 if self.empty(): 

61 raise ValueError("heap is empty") 

62 return heapq.nsmallest(1, self)[0] 

63 

64 def push(self, elt: T) -> Optional[T]: 

65 """Pushes a new element onto the heap. 

66 Returns the element that was removed, if one exists.""" 

67 if (self.N is None) or (len(self) < self.N): 

68 heapq.heappush(self, elt) 

69 return None 

70 return heapq.heappushpop(self, elt) 

71 

72 def pop(self) -> T: # type: ignore[override] 

73 """Pops off the smallest element from the heap and returns it.""" 

74 return heapq.heappop(self)