Skip to content

Commit f8b5e87

Browse files
QubitiumQubitium
committed
refractor set_up_logger with LobBar.shared() api
Signed-off-by: Qubitium <qubitium@modelcloud.ai>
1 parent 299da3e commit f8b5e87

File tree

4 files changed

+198
-179
lines changed

4 files changed

+198
-179
lines changed

‎logbar/logbar.py‎

Lines changed: 173 additions & 156 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import logging
1818
from enum import Enum
19+
from typing import Optional, Iterable
1920

2021
from logbar.terminal import terminal_size
2122

@@ -44,171 +45,187 @@ class LEVEL(str, Enum):
4445
ERROR = "ERROR"
4546
CRITICAL = "CRIT"
4647

47-
def setup_logger():
48-
global logger
49-
if logger is not None:
48+
class LogBar(logging.Logger):
49+
history = set()
50+
history_limit = 1000
51+
52+
@classmethod
53+
# return a shared global/singleton logger
54+
def shared(cls, override_logger: Optional[bool] = False):
55+
global logger
56+
if logger is not None:
57+
return logger
58+
59+
# save logger class
60+
if not override_logger:
61+
original_logger_cls = logging.getLoggerClass()
62+
63+
logging.setLoggerClass(LogBar)
64+
65+
logger = logging.getLogger("logbar")
66+
67+
# restore logger cls
68+
if not override_logger:
69+
logging.setLoggerClass(original_logger_cls)
70+
71+
logger.propagate = False
72+
logger.setLevel(logging.INFO)
73+
74+
# handler = logging.StreamHandler(sys.stdout)
75+
# handler.setFormatter(formatter)
76+
# handler.flush = sys.stdout.flush
77+
# logger.addHandler(handler)
78+
79+
# clear space from previous logs
80+
print("", end='\n', flush=True)
81+
5082
return logger
5183

52-
class CustomLogger(logging.Logger):
53-
history = set()
54-
history_limit = 1000
5584

56-
def history_add(self, msg) -> bool:
57-
h = hash(msg) # TODO only msg is checked not level + msg
58-
if h in self.history:
59-
return False # add failed since it already exists
85+
def pb(self, iterable: Iterable):
86+
from logbar.progress import ProgressBar
87+
88+
return ProgressBar(iterable)
89+
90+
def history_add(self, msg) -> bool:
91+
h = hash(msg) # TODO only msg is checked not level + msg
92+
if h in self.history:
93+
return False # add failed since it already exists
94+
95+
if len(self.history) > self.history_limit:
96+
self.history.clear()
97+
98+
self.history.add(h)
99+
100+
return True
101+
102+
class critical_cls:
103+
def __init__(self, logger):
104+
self.logger = logger
105+
106+
def once(self, msg, *args, **kwargs):
107+
if self.logger.history_add(msg):
108+
self(msg, *args, **kwargs)
109+
110+
def __call__(self, msg, *args, **kwargs):
111+
self.logger._process(LEVEL.CRITICAL, msg, *args, **kwargs)
112+
113+
class warn_cls:
114+
def __init__(self, logger):
115+
self.logger = logger
116+
117+
def once(self, msg, *args, **kwargs):
118+
if self.logger.history_add(msg):
119+
self(msg, *args, **kwargs)
120+
121+
def __call__(self, msg, *args, **kwargs):
122+
self.logger._process(LEVEL.WARN, msg, *args, **kwargs)
123+
124+
class debug_cls:
125+
def __init__(self, logger):
126+
self.logger = logger
127+
128+
def once(self, msg, *args, **kwargs):
129+
if self.logger.history_add(msg):
130+
self(msg, *args, **kwargs)
131+
132+
def __call__(self, msg, *args, **kwargs):
133+
self.logger._process(LEVEL.DEBUG, msg, *args, **kwargs)
134+
135+
class info_cls:
136+
def __init__(self, logger):
137+
self.logger = logger
138+
139+
def once(self, msg, *args, **kwargs):
140+
if self.logger.history_add(msg):
141+
self(msg, *args, **kwargs)
142+
143+
def __call__(self, msg, *args, **kwargs):
144+
self.logger._process(LEVEL.INFO, msg, *args, **kwargs)
145+
146+
class error_cls:
147+
def __init__(self, logger):
148+
self.logger = logger
149+
150+
def once(self, msg, *args, **kwargs):
151+
if self.logger.history_add(msg):
152+
self(msg, *args, **kwargs)
153+
154+
def __call__(self, msg, *args, **kwargs):
155+
self.logger._process(LEVEL.ERROR, msg, *args, **kwargs)
156+
157+
def __init__(self, name):
158+
super().__init__(name)
159+
self._warning = self.warning
160+
self._debug = self.debug
161+
self._info = self.info
162+
self._error = self.error
163+
self._critical = self.critical
164+
165+
self.warn = self.warn_cls(logger=self)
166+
self.debug = self.debug_cls(logger=self)
167+
self.info = self.info_cls(logger=self)
168+
self.error = self.error_cls(logger=self)
169+
self.critical = self.critical_cls(logger=self)
60170

61-
if len(self.history) > self.history_limit:
62-
self.history.clear()
171+
def _process(self, level: LEVEL, msg, *args, **kwargs):
172+
from logbar.progress import ProgressBar
63173

64-
self.history.add(h)
174+
columns, _ = terminal_size()
175+
str_msg = str(msg)
65176

66-
return True
177+
global last_pb_instance
178+
if isinstance(last_pb_instance, ProgressBar) and not last_pb_instance.closed:
179+
buf = f'\r'
180+
if columns > 0:
181+
str_msg += " " * columns
67182

68-
class critical_cls:
69-
def __init__(self, logger):
70-
self.logger = logger
183+
print(buf,end='',flush=True)
71184

72-
def once(self, msg, *args, **kwargs):
73-
if self.logger.history_add(msg):
74-
self(msg, *args, **kwargs)
185+
# Get the color for the log level
186+
187+
reset = COLORS["RESET"]
188+
color = COLORS.get(level.value, reset)
189+
190+
out_len = 5 + 1 + len(str_msg)
191+
paddding_end = " " * (columns - out_len)
192+
193+
padding = " " * (5 - len(level.value)) # 5 is max enum string length
194+
print(f"{color}{level.value}{reset}{padding} {str_msg}", end='\n', flush=True)
195+
# if level == LEVEL.INFO:
196+
# print(f"INFO: {str_msg}", end='',flush=True)
197+
# elif level == LEVEL.WARN:
198+
# print(f"WARN: {str_msg}", end='',flush=True)
199+
# elif level == LEVEL.ERROR:
200+
# print(f"ERROR: {str_msg}", end='',flush=True)
201+
# elif level == LEVEL.DEBUG:
202+
# print(f"DEBUG: {str_msg}", end='',flush=True)
203+
# else:
204+
# raise RuntimeError(f"Unknown logging level {level}")
205+
206+
# Print the message with the appropriate color
207+
#print(f"{color}{level.value}{reset}{padding} {str_msg}")
208+
209+
# if level == LEVEL.INFO:
210+
# self._info(str_msg, *args, **kwargs)
211+
#
212+
# elif level == LEVEL.WARN:
213+
# self._warning(str_msg, *args, **kwargs)
214+
# elif level == LEVEL.ERROR:
215+
# self._error(str_msg, *args, **kwargs)
216+
# elif level == LEVEL.DEBUG:
217+
# self._debug(str_msg, *args, **kwargs)
218+
# else:
219+
# raise RuntimeError(f"Unknown logging level {level}")
220+
221+
if isinstance(last_pb_instance, ProgressBar):
222+
if not last_pb_instance.closed:
223+
# only do this for our instance
224+
if self == logger:
225+
last_pb_instance.draw()
226+
else:
227+
last_pb_instance = None
75228

76-
def __call__(self, msg, *args, **kwargs):
77-
self.logger._process(LEVEL.CRITICAL, msg, *args, **kwargs)
78229

79-
class warn_cls:
80-
def __init__(self, logger):
81-
self.logger = logger
82-
83-
def once(self, msg, *args, **kwargs):
84-
if self.logger.history_add(msg):
85-
self(msg, *args, **kwargs)
86-
87-
def __call__(self, msg, *args, **kwargs):
88-
self.logger._process(LEVEL.WARN, msg, *args, **kwargs)
89-
90-
class debug_cls:
91-
def __init__(self, logger):
92-
self.logger = logger
93-
94-
def once(self, msg, *args, **kwargs):
95-
if self.logger.history_add(msg):
96-
self(msg, *args, **kwargs)
97-
98-
def __call__(self, msg, *args, **kwargs):
99-
self.logger._process(LEVEL.DEBUG, msg, *args, **kwargs)
100-
101-
class info_cls:
102-
def __init__(self, logger):
103-
self.logger = logger
104-
105-
def once(self, msg, *args, **kwargs):
106-
if self.logger.history_add(msg):
107-
self(msg, *args, **kwargs)
108-
109-
def __call__(self, msg, *args, **kwargs):
110-
self.logger._process(LEVEL.INFO, msg, *args, **kwargs)
111-
112-
class error_cls:
113-
def __init__(self, logger):
114-
self.logger = logger
115-
116-
def once(self, msg, *args, **kwargs):
117-
if self.logger.history_add(msg):
118-
self(msg, *args, **kwargs)
119-
120-
def __call__(self, msg, *args, **kwargs):
121-
self.logger._process(LEVEL.ERROR, msg, *args, **kwargs)
122-
123-
def __init__(self, name):
124-
super().__init__(name)
125-
self._warning = self.warning
126-
self._debug = self.debug
127-
self._info = self.info
128-
self._error = self.error
129-
self._critical = self.critical
130-
131-
self.warn = self.warn_cls(logger=self)
132-
self.debug = self.debug_cls(logger=self)
133-
self.info = self.info_cls(logger=self)
134-
self.error = self.error_cls(logger=self)
135-
self.critical = self.critical_cls(logger=self)
136-
137-
def _process(self, level: LEVEL, msg, *args, **kwargs):
138-
from logbar.progress import ProgressBar
139-
140-
columns, _ = terminal_size()
141-
str_msg = str(msg)
142-
143-
global last_pb_instance
144-
if isinstance(last_pb_instance, ProgressBar) and not last_pb_instance.closed:
145-
buf = f'\r'
146-
if columns > 0:
147-
str_msg += " " * columns
148-
149-
print(buf,end='',flush=True)
150-
151-
# Get the color for the log level
152-
153-
reset = COLORS["RESET"]
154-
color = COLORS.get(level.value, reset)
155-
156-
out_len = 5 + 1 + len(str_msg)
157-
paddding_end = " " * (columns - out_len)
158-
159-
padding = " " * (5 - len(level.value)) # 5 is max enum string length
160-
print(f"{color}{level.value}{reset}{padding} {str_msg}", end='\n', flush=True)
161-
# if level == LEVEL.INFO:
162-
# print(f"INFO: {str_msg}", end='',flush=True)
163-
# elif level == LEVEL.WARN:
164-
# print(f"WARN: {str_msg}", end='',flush=True)
165-
# elif level == LEVEL.ERROR:
166-
# print(f"ERROR: {str_msg}", end='',flush=True)
167-
# elif level == LEVEL.DEBUG:
168-
# print(f"DEBUG: {str_msg}", end='',flush=True)
169-
# else:
170-
# raise RuntimeError(f"Unknown logging level {level}")
171-
172-
# Print the message with the appropriate color
173-
#print(f"{color}{level.value}{reset}{padding} {str_msg}")
174-
175-
# if level == LEVEL.INFO:
176-
# self._info(str_msg, *args, **kwargs)
177-
#
178-
# elif level == LEVEL.WARN:
179-
# self._warning(str_msg, *args, **kwargs)
180-
# elif level == LEVEL.ERROR:
181-
# self._error(str_msg, *args, **kwargs)
182-
# elif level == LEVEL.DEBUG:
183-
# self._debug(str_msg, *args, **kwargs)
184-
# else:
185-
# raise RuntimeError(f"Unknown logging level {level}")
186-
187-
if isinstance(last_pb_instance, ProgressBar):
188-
if not last_pb_instance.closed:
189-
# only do this for our instance
190-
if self == logger:
191-
last_pb_instance.draw()
192-
else:
193-
last_pb_instance = None
194-
195-
# original_logger_cls = logging.getLoggerClass()
196-
logging.setLoggerClass(CustomLogger)
197-
198-
logger = logging.getLogger("logbar")
199-
# logging.setLoggerClass(original_logger_cls)
200-
201-
logger.propagate = False
202-
logger.setLevel(logging.INFO)
203-
204-
# handler = logging.StreamHandler(sys.stdout)
205-
# handler.setFormatter(formatter)
206-
# handler.flush = sys.stdout.flush
207-
# logger.addHandler(handler)
208-
209-
# clear space from previous logs
210-
print("", end='\n', flush=True)
211-
212-
return logger
213230

214231

‎logbar/progress.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@
2020
from typing import Iterable, Optional
2121
from warnings import warn
2222

23-
from logbar.logbar import setup_logger, update_last_pb_instance
23+
from logbar.logbar import update_last_pb_instance, LogBar
2424
from logbar.terminal import terminal_size
2525

26-
logger = setup_logger()
26+
logger = LogBar.shared()
2727

2828
# TODO FIXME: what does this do exactly?
2929
class ProgressBarWarning(Warning):

‎tests/test_log.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,21 @@
11
import unittest
22

3-
from logbar.logbar import setup_logger
3+
from logbar.logbar import LogBar
44

5-
logger = setup_logger()
5+
log = LogBar.shared()
66

77
class TestProgressBar(unittest.TestCase):
88

99
def test_log_simple(self):
10-
logger.info("hello info")
10+
log.info("hello info")
1111

1212
def test_log_once(self):
13-
logger.info.once("hello info")
14-
logger.info.once("hello info")
13+
log.info.once("hello info 1")
14+
log.info.once("hello info 1")
1515

1616
def test_levels(self):
17-
logger.info("hello info")
18-
logger.debug("hello debug")
19-
logger.warn("hello warn")
20-
logger.error("hello error")
21-
logger.critical("hello critical")
17+
log.info("hello info")
18+
log.debug("hello debug")
19+
log.warn("hello warn")
20+
log.error("hello error")
21+
log.critical("hello critical")

0 commit comments

Comments
 (0)