chore(//py/torch_tensorrt): conforming logging to mypy

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

Author: narendasan

py/torch_tensorrt/logging.py

@@ -1,3 +1,4 @@
+from typing import Self, Any
 from enum import Enum
 from torch_tensorrt._C import (
     _get_logging_prefix,
@@ -22,19 +23,21 @@ class Level(Enum):
     Graph = LogLevel.GRAPH
 
     @staticmethod
-    def _to_internal_level(external) -> LogLevel:
+    def _to_internal_level(external: Self) -> LogLevel:
         if external == Level.InternalError:
             return LogLevel.INTERNAL_ERROR
-        if external == Level.Error:
+        elif external == Level.Error:
             return LogLevel.ERROR
-        if external == Level.Warning:
+        elif external == Level.Warning:
             return LogLevel.WARNING
-        if external == Level.Info:
+        elif external == Level.Info:
             return LogLevel.INFO
-        if external == Level.Debug:
+        elif external == Level.Debug:
             return LogLevel.DEBUG
-        if external == Level.Graph:
+        elif external == Level.Graph:
             return LogLevel.GRAPH
+        else:
+            raise ValueError("Unknown log severity")
 
 
 def get_logging_prefix() -> str:
@@ -43,10 +46,10 @@ def get_logging_prefix() -> str:
     Returns:
         str: Prefix used for logger
     """
-    return _get_logging_prefix()
+    return str(_get_logging_prefix())
 
 
-def set_logging_prefix(prefix: str):
+def set_logging_prefix(prefix: str) -> None:
     """Set the prefix used when logging messages
 
     Args:
@@ -64,7 +67,7 @@ def get_reportable_log_level() -> Level:
     return Level(_get_reportable_log_level())
 
 
-def set_reportable_log_level(level: Level):
+def set_reportable_log_level(level: Level) -> None:
     """Set the level required for a message to be printed to the log
 
     Args:
@@ -79,10 +82,10 @@ def get_is_colored_output_on() -> bool:
     Returns:
         bool: If colored output is one
     """
-    return _get_is_colored_output_on()
+    return bool(_get_is_colored_output_on())
 
 
-def set_is_colored_output_on(colored_output_on: bool):
+def set_is_colored_output_on(colored_output_on: bool) -> None:
     """Enable or disable color in the log output
 
     Args:
@@ -91,7 +94,7 @@ def set_is_colored_output_on(colored_output_on: bool):
     _set_is_colored_output_on(colored_output_on)
 
 
-def log(level: Level, msg: str):
+def log(level: Level, msg: str) -> None:
     """Add a new message to the log
 
     Adds a new message to the log at a specified level. The message
@@ -120,11 +123,11 @@ class internal_errors:
         outputs = model_torchtrt(inputs)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.InternalError)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)
 
 
@@ -137,11 +140,11 @@ class errors:
         outputs = model_torchtrt(inputs)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Error)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)
 
 
@@ -154,11 +157,11 @@ class warnings:
         model_trt = torch_tensorrt.compile(model, **spec)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Warning)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)
 
 
@@ -171,11 +174,11 @@ class info:
         model_trt = torch_tensorrt.compile(model, **spec)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Info)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)
 
 
@@ -188,11 +191,11 @@ class debug:
         model_trt = torch_tensorrt.compile(model, **spec)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Debug)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)
 
 
@@ -206,9 +209,9 @@ class graphs:
         model_trt = torch_tensorrt.compile(model, **spec)
     """
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Graph)
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
         set_reportable_log_level(self.external_lvl)