chore(//py): Adding documentation to the new context managers

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

Author: narendasan

py/torch_tensorrt/logging.py

@@ -99,55 +99,109 @@ def log(level: Level, msg: str):
     """
     _log(Level._to_internal_level(level), msg)
 
-
     InternalError = LogLevel.INTERNAL_ERROR
     Error = LogLevel.ERROR
     Warning = LogLevel.WARNING
     Info = LogLevel.INFO
     Debug = LogLevel.DEBUG
     Graph = LogLevel.GRAPH
 
+
 class internal_errors:
+    """Context-manager to limit displayed log messages to just internal errors
+
+    Example::
+
+    with torch_tensorrt.logging.internal_errors():
+        outputs = model_torchtrt(inputs)
+    """
+
     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:
+    """Context-manager to limit displayed log messages to just errors and above
+
+    Example::
+
+    with torch_tensorrt.logging.errors():
+        outputs = model_torchtrt(inputs)
+    """
+
     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:
+    """Context-manager to limit displayed log messages to just warnings and above
+
+    Example::
+
+    with torch_tensorrt.logging.warnings():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
     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:
+    """Context-manager to display all info and greater severity messages
+
+    Example::
+
+    with torch_tensorrt.logging.info():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
     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:
+    """Context-manager to display full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.debug():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
     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:
+    """Context-manager to display the results of intermediate lowering passes
+    as well as full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.graphs():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Graph)