ANL-CEEESA/MIPLearn
1
0
Fork 0
mirror of https://github.com/ANL-CEEESA/MIPLearn.git synced 2025-12-06 09:28:51 -06:00
Code Issues Packages Projects Releases Wiki Activity

Add types to log.py

Browse Source
This commit is contained in:
Alinson S. Xavier
2021-04-07 21:01:21 -05:00
parent 331ee5914d
commit f5606efb72
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
miplearn/log.py
View File

@@ -7,17 +7,22 @@ import sys
import time import time
import traceback import traceback
import warnings import warnings
from typing import Dict, Any, Optional
_formatwarning = warnings.formatwarning _formatwarning = warnings.formatwarning
class TimeFormatter(logging.Formatter): class TimeFormatter(logging.Formatter):
def __init__(self, start_time, log_colors): def __init__(
self,
start_time: float,
log_colors: Dict[str, str],
) -> None:
super().__init__() super().__init__()
self.start_time = start_time self.start_time = start_time
self.log_colors = log_colors self.log_colors = log_colors
def format(self, record): def format(self, record: logging.LogRecord) -> str:
if record.levelno >= logging.ERROR: if record.levelno >= logging.ERROR:
color = self.log_colors["red"] color = self.log_colors["red"]
elif record.levelno >= logging.WARNING: elif record.levelno >= logging.WARNING:
@@ -32,14 +37,17 @@ class TimeFormatter(logging.Formatter):
) )
def formatwarning_tb(*args, **kwargs): def formatwarning_tb(*args: Any, **kwargs: Any) -> str:
s = _formatwarning(*args, **kwargs) s = _formatwarning(*args, **kwargs)
tb = traceback.format_stack() tb = traceback.format_stack()
s += "".join(tb[:-1]) s += "".join(tb[:-1])
return s return s
def setup_logger(start_time=None, force_color=False): def setup_logger(
start_time: Optional[float] = None,
force_color: bool = False,
) -> None:
if start_time is None: if start_time is None:
start_time = time.time() start_time = time.time()
if sys.stdout.isatty() or force_color: if sys.stdout.isatty() or force_color: