another implementation of log_interval

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/543

The previous implementation:
> the problem is the ContextDecorator somehow swallows the exception in the wrapped function and just returns None.

This diff adds a test such that previous implementation would fail:
```
======================================================================
FAIL: test_log_interval_error_prop (d2go.tests.fb.test_utils_logging.TestUtilsLogging)
Make sure the log_interval can handle error propagation.
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/sandcastle/boxes/fbsource/buck-out/v2/gen/fbcode/ef4169ac7f95fb74/mobile-vision/d2go/tests/__init_tests__/init_tests#link-tree/d2go/tests/fb/test_utils_logging.py", line 152, in test_log_interval_error_prop
    foo(-1)
AssertionError: ValueError not raised

----------------------------------------------------------------------
Ran 1 test in 0.098s
```

The new version seems easier to understand and doesn't have the error swallowing.

Reviewed By: jaconey

Differential Revision: D46009938

fbshipit-source-id: 6b632deb513ab47c4d760f796bf49fc45eae3005

d2go/utils/logging.py

@@ -5,8 +5,8 @@ import logging
 import sys
 import time
 import uuid
-from contextlib import ContextDecorator
-from typing import Any, Optional
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar
 
 from mobile_cv.common.misc.oss_utils import fb_overwritable
 
@@ -14,6 +14,8 @@ from mobile_cv.common.misc.oss_utils import fb_overwritable
 # Saving the builtin print to wrap it up later.
 BUILTIN_PRINT = builtins.print
 
+_T = TypeVar("_T")
+
 
 @fb_overwritable()
 def initialize_logging(logging_level: int) -> None:
@@ -64,34 +66,27 @@ def _log_exit(category: str, name: str, unique_id: str, duration: float) -> None
     )
 
 
-class log_interval(ContextDecorator):
-    def __init__(
-        self, category: Optional[str] = None, name: Optional[str] = None
-    ) -> None:
-        super().__init__()
-        self._unique_id = uuid.uuid1().int >> 97
-        self._category = category
-        self._name = name
-        self._start = 0
-
-    def __call__(self, func):
-        if self._category is None:
-            self._category = func.__qualname__.split(".")[0]
-        if self._name is None:
-            self._name = func.__name__
-
-        return super().__call__(func)
-
-    def __enter__(self) -> "log_interval":
-        _log_enter(self._category, self._name, self._unique_id)
-        self._start = time.perf_counter()
-        return self
-
-    def __exit__(self, exc_type, exc_value, tb) -> bool:
-        _log_exit(
-            self._category,
-            self._name,
-            self._unique_id,
-            time.perf_counter() - self._start,
-        )
-        return True
+def log_interval(
+    category: Optional[str] = None, name: Optional[str] = None
+) -> Callable[[Callable[..., _T]], Callable[..., _T]]:
+
+    _unique_id = uuid.uuid1().int >> 97
+    _overwrite_category = category
+    _overwrite_name = name
+
+    def log_interval_deco(func: Callable[..., _T]) -> Callable[..., _T]:
+
+        _category = _overwrite_category or func.__qualname__.split(".")[0]
+        _name = _overwrite_name or func.__name__
+
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> _T:
+            _log_enter(_category, _name, _unique_id)
+            _start = time.perf_counter()
+            ret = func(*args, **kwargs)
+            _log_exit(_category, _name, _unique_id, time.perf_counter() - _start)
+            return ret
+
+        return wrapper
+
+    return log_interval_deco