refactor(//py): Give better names to the CtxManagers

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/torch_tensorrt/logging.py

@@ -107,51 +107,50 @@ def log(level: Level, msg: str):
     Debug = LogLevel.DEBUG
     Graph = LogLevel.GRAPH
 
-class InternalErrors:
+class internal_errors:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.InternalError)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
-class Errors:
+class errors:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Error)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
-class Warnings:
+class warnings:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Warning)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
-class Info:
+class info:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Info)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
-class Debug:
+class debug:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Debug)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
-class Graphs:
+class graphs:
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Graph)
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
-

tests/py/test_api.py

@@ -323,6 +323,49 @@ def test_is_colored_output_on(self):
         color = torchtrt.logging.get_is_colored_output_on()
         self.assertTrue(color)
 
+    def test_context_managers(self):
+        base_lvl = torchtrt.logging.get_reportable_log_level()
+        with torchtrt.logging.internal_errors():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.InternalError, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.errors():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Error, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.warnings():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Warning, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.info():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Info, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.debug():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Debug, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.graphs():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Graph, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
 
 class TestDevice(unittest.TestCase):