Coverage for optimates/utils.py: 59%
35 statements
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-16 11:19 -0500
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-16 11:19 -0500
1"""Utility module."""
3import heapq
4import logging
5from typing import Any, Optional, TypeVar, cast
8T = TypeVar('T')
11###########
12# LOGGING #
13###########
15class Logger(logging.Logger):
16 """Custom subclass of logging.Logger."""
18 def __init__(self, *args: Any, **kwargs: Any) -> None:
19 super().__init__(*args, **kwargs)
20 self.verbosity = 0
22 def level_for_verbosity(self, verbosity: int) -> int:
23 """Converts a verbosity level to a logging level."""
24 return logging.INFO - verbosity
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))
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)
37LOG_FMT = '%(name)s - %(message)s'
38logging.basicConfig(format=LOG_FMT, level=logging.INFO)
39logging.setLoggerClass(Logger)
40logger = cast(Logger, logging.getLogger('optimates'))
43########
44# HEAP #
45########
47class TopNHeap(list[T]):
48 """Maintains the largest N elements on a heap."""
50 def __init__(self, N: Optional[int] = None) -> None:
51 super().__init__()
52 self.N = N
54 def empty(self) -> bool:
55 """Returns True if the heap is empty."""
56 return (len(self) == 0)
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]
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)
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)