Reduce complexity of logging try-catch + add typing (pytorch#1380)

Summary:
Pull Request resolved: pytorch#1380

'TryExcept 6' is too complex (13) See https://www.flake8rules.com/rules/C901.html

Reviewed By: jjuncho

Differential Revision: D64511308

fbshipit-source-id: c42173f6e7ea4b8732141640f4886b3e347f599a

captum/log/__init__.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 # pyre-strict
-from typing import Any
 
 try:
     from captum.log.fb.internal_log import (
@@ -19,49 +18,16 @@
         "TimedLog",
         "set_environment",
         "disable_detailed_logging",
+        "patch_methods",
     ]
 
 except ImportError:
-    from functools import wraps
-
-    def log(*args: Any, **kwargs: Any) -> None:  # type: ignore
-        pass
-
     # bug with mypy: https://github.com/python/mypy/issues/1153
-    class TimedLog:  # type: ignore
-        # pyre-fixme[2]: Parameter must be annotated.
-        def __init__(self, *args, **kwargs) -> None:
-            pass
-
-        def __enter__(self) -> "TimedLog":
-            return self
-
-        # pyre-fixme[2]: Parameter must be annotated.
-        def __exit__(self, exception_type, exception_value, traceback) -> bool:
-            return exception_value is not None
-
-    # pyre-fixme[3]: Return type must be annotated.
-    def log_usage(*log_args: Any, **log_kwargs: Any):
-        # pyre-fixme[3]: Return type must be annotated.
-        # pyre-fixme[2]: Parameter must be annotated.
-        def _log_usage(func):
-            @wraps(func)
-            # pyre-fixme[53]: Captured variable `func` is not annotated.
-            # pyre-fixme[3]: Return type must be annotated.
-            def wrapper(*args: Any, **kwargs: Any):
-                return func(*args, **kwargs)
-
-            return wrapper
-
-        return _log_usage
-
-    # pyre-fixme[2]: Parameter must be annotated.
-    def set_environment(env) -> None:  # type: ignore
-        pass
-
-    def disable_detailed_logging() -> None:
-        pass
-
-    # pyre-fixme[2]: Parameter must be annotated.
-    def patch_methods(tester, patch_log: bool = True) -> None:  # type: ignore
-        pass
+    from captum.log.dummy_log import (  # type: ignore
+        disable_detailed_logging,
+        log,
+        log_usage,
+        patch_methods,
+        set_environment,
+        TimedLog,
+    )

captum/log/dummy_log.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+
+# pyre-strict
+
+from functools import wraps
+from types import TracebackType
+from typing import Any, List, Optional, Union
+
+
+def log(*args: Any, **kwargs: Any) -> None:
+    pass
+
+
+class TimedLog:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def __enter__(self) -> "TimedLog":
+        return self
+
+    def __exit__(
+        self,
+        exception_type: Optional[BaseException],
+        exception_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        return exception_value is not None
+
+
+# pyre-fixme[3]: Return type must be annotated.
+def log_usage(*log_args: Any, **log_kwargs: Any):
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
+    def _log_usage(func):
+        @wraps(func)
+        # pyre-fixme[53]: Captured variable `func` is not annotated.
+        # pyre-fixme[3]: Return type must be annotated.
+        def wrapper(*args: Any, **kwargs: Any):
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return _log_usage
+
+
+def set_environment(env: Union[None, List[str], str]) -> None:
+    pass
+
+
+def disable_detailed_logging() -> None:
+    pass
+
+
+def patch_methods(_, patch_log: bool = True) -> None:
+    pass