Skip to content

Commit 12e470f

Browse files
committed
feat(//py): Context managers to quickly switch logging level
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent ef62f6b commit 12e470f

File tree

1 file changed

+59
-0
lines changed

1 file changed

+59
-0
lines changed

py/torch_tensorrt/logging.py

+59
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
from enum import Enum
2+
from hashlib import new
3+
from imp import new_module
24
from torch_tensorrt._C import _get_logging_prefix, _set_logging_prefix, \
35
_get_reportable_log_level, _set_reportable_log_level, \
46
_get_is_colored_output_on, _set_is_colored_output_on, \
@@ -96,3 +98,60 @@ def log(level: Level, msg: str):
9698
msg (str): Actual message text
9799
"""
98100
_log(Level._to_internal_level(level), msg)
101+
102+
103+
InternalError = LogLevel.INTERNAL_ERROR
104+
Error = LogLevel.ERROR
105+
Warning = LogLevel.WARNING
106+
Info = LogLevel.INFO
107+
Debug = LogLevel.DEBUG
108+
Graph = LogLevel.GRAPH
109+
110+
class InternalErrors:
111+
def __enter__(self):
112+
self.external_lvl = get_reportable_log_level()
113+
set_reportable_log_level(Level.InternalError)
114+
115+
def __exit__(self, exc_type, exc_value, exc_tb):
116+
set_reportable_log_level(self.external_lvl)
117+
118+
class Errors:
119+
def __enter__(self):
120+
self.external_lvl = get_reportable_log_level()
121+
set_reportable_log_level(Level.Error)
122+
123+
def __exit__(self, exc_type, exc_value, exc_tb):
124+
set_reportable_log_level(self.external_lvl)
125+
126+
class Warnings:
127+
def __enter__(self):
128+
self.external_lvl = get_reportable_log_level()
129+
set_reportable_log_level(Level.Warning)
130+
131+
def __exit__(self, exc_type, exc_value, exc_tb):
132+
set_reportable_log_level(self.external_lvl)
133+
134+
class Info:
135+
def __enter__(self):
136+
self.external_lvl = get_reportable_log_level()
137+
set_reportable_log_level(Level.Info)
138+
139+
def __exit__(self, exc_type, exc_value, exc_tb):
140+
set_reportable_log_level(self.external_lvl)
141+
142+
class Debug:
143+
def __enter__(self):
144+
self.external_lvl = get_reportable_log_level()
145+
set_reportable_log_level(Level.Debug)
146+
147+
def __exit__(self, exc_type, exc_value, exc_tb):
148+
set_reportable_log_level(self.external_lvl)
149+
150+
class Graphs:
151+
def __enter__(self):
152+
self.external_lvl = get_reportable_log_level()
153+
set_reportable_log_level(Level.Graph)
154+
155+
def __exit__(self, exc_type, exc_value, exc_tb):
156+
set_reportable_log_level(self.external_lvl)
157+

0 commit comments

Comments
 (0)