Add a simple util function

src/hello_world.py

@@ -2,6 +2,8 @@
 import datetime
 from typing import List, Optional
 
+import src.utils
+
 
 @dataclasses.dataclass
 class Person:
@@ -19,6 +21,7 @@ def find_age(
     return result[0].age
 
 
+@src.utils.timer(lambda t: print("Function add takes " + str(t) + "ns"))
 def add(a: int, b: int) -> int:
     return a + b
 

src/utils.py

@@ -0,0 +1,22 @@
+import time
+from typing import Any, Callable, TypeVar
+
+RT = TypeVar("RT")  # return type
+
+
+def timer(
+    output_time_func: Callable[[int], None],
+) -> Callable[[Callable[..., RT]], Callable[..., RT]]:
+    """Decorator to measure performance"""
+
+    def decorator(func: Callable[..., RT]) -> Callable[..., RT]:
+        def wrapper(*args: Any, **kwargs: Any) -> RT:
+            t1 = time.perf_counter_ns()
+            result = func(*args, **kwargs)
+            t2 = time.perf_counter_ns()
+            output_time_func(t2 - t1)
+            return result
+
+        return wrapper
+
+    return decorator

tests/unit/test_utils.py

@@ -0,0 +1,31 @@
+import time
+import unittest
+from typing import List
+
+import src.utils
+
+
+class TestBasic(unittest.TestCase):
+    def test_timer_records_correct_time(self) -> None:
+        @src.utils.timer(lambda t: self.assertGreater(t, 1e9))
+        def wait_one_sec() -> None:
+            time.sleep(1)
+
+        self.assertIsNone(wait_one_sec())
+
+    def test_timer_decorator_returns_correct_output(self) -> None:
+        @src.utils.timer(lambda _: None)
+        def return_int() -> int:
+            return 13
+
+        @src.utils.timer(lambda _: None)
+        def return_str() -> str:
+            return "hello world"
+
+        @src.utils.timer(lambda _: None)
+        def return_list() -> List[int]:
+            return [1, 2, 3]
+
+        self.assertEqual(return_int(), 13)
+        self.assertEqual(return_str(), "hello world")
+        self.assertListEqual(return_list(), [1, 2, 3])