feat: add timing information to ic debugger

icecream/icecream.py

@@ -16,7 +16,8 @@
 import inspect
 import pprint
 import sys
-from types import FrameType
+import time
+from types import FrameType, TracebackType
 from typing import (
     Optional,
     cast,
@@ -361,7 +362,6 @@ def _formatArgs(
         context: str,
         args: Sequence[object]
     ) -> str:
-
         callNode = Source.executing(callFrame).node
         if callNode is not None:
             assert isinstance(callNode, ast.Call)
@@ -530,5 +530,62 @@ def configureOutput(
         if lineWrapWidth is not Sentinel.absent:
             self.lineWrapWidth = lineWrapWidth
 
+    @property
+    def timer(self) -> "Timer":
+        return Timer(self)
+
+
+class Timer:
+    def __init__(self, ic: IceCreamDebugger):
+        self._ic = ic
+        self._enter_time: Optional[float] = None
+
+    def format_duration(self, seconds: float) -> str:
+        if seconds < 1e-6:
+            return f"{seconds * 1e9:.2f}ns"
+        if seconds < 1e-3:
+            return f"{seconds * 1e6:.2f}us"
+        if seconds < 1:
+            return f"{seconds * 1e3:.2f}ms"
+        if seconds < 60:
+            return f"{seconds:.2f}s"
+        if seconds < 3600:
+            return f"{int(seconds // 60)}m {seconds % 60:.2f}s"
+        return f"{int(seconds // 3600)}h {int((seconds % 3600) // 60)}m {seconds % 60:.2f}s"
+
+    def _output(self, duration: float, label: str = "") -> None:
+        if not self._ic.enabled:
+            return
+        prefix = cast(str, call_or_value(self._ic.prefix))
+        formatted_time = self.format_duration(duration)
+        if label:
+            msg = f"{prefix}{label} took {formatted_time}"
+        else:
+            msg = f"{formatted_time}"
+        self._ic.outputFunction(msg)
+
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            start_time: float = time.perf_counter()
+            try:
+                return func(*args, **kwargs)
+            finally:
+                duration: float = time.perf_counter() - start_time
+                self._output(duration, func.__name__)
+
+        return wrapper
+
+    def __enter__(self) -> "Timer":
+        self._enter_time = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType]) -> None:
+        if self._enter_time is None:
+            raise RuntimeError("Timer.__exit__ called without __enter__. ")
+        duration: float = time.perf_counter() - self._enter_time
+        self._output(duration)
+        self._enter_time = None
+
 
 ic = IceCreamDebugger()

tests/test_icecream.py

@@ -9,8 +9,9 @@
 #
 # License: MIT
 #
-
+import re
 import sys
+import time
 import unittest
 import warnings
 
@@ -24,6 +25,8 @@
 from icecream.icecream import has_non_ascii_chars
 
 TEST_PAIR_DELIMITER = '| '
+TIMER_DURATION_RE_DECORATOR = r'\d+\.\d{2}(ns|us|ms|s)'
+TIMER_DURATION_RE_CONTEXT_MANAGER = r'(\d+\.\d{2})(ns|us|ms|s)'
 MY_FILENAME = basename(__file__)
 MY_FILEPATH = realpath(__file__)
 
@@ -37,6 +40,16 @@ def noop(*args, **kwargs):  # type: ignore
     return
 
 
+@ic.timer
+def noop_with_time_decorator(*args, **kwargs):  # type: ignore
+    return
+
+
+@ic.timer
+def raise_value_error():
+    raise ValueError('test error')
+
+
 def has_ansi_escape_codes(s: str) -> bool:
     # oversimplified, but ¯\_(ツ)_/¯. TODO(grun): Test with regex.
     return '\x1b[' in s
@@ -798,3 +811,165 @@ def test_no_color_with_explicit_output_function(self):
         finally:
             ic.configureOutput(noColor=originalNoColor)
             ic.outputFunction = originalOutputFunction
+
+    def test_with_timer_decorator(self):
+        @ic.timer
+        def timed_noop():
+            pass
+
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            timed_noop()
+
+        output = err.getvalue().strip()
+        func_name = timed_noop.__name__
+
+        pattern = rf"^ic\| {func_name} took {TIMER_DURATION_RE_DECORATOR}$"
+        self.assertRegex(output, pattern)
+
+    def test_timer_decorator_nested(self):
+        @ic.timer
+        def outer_method():
+            @ic.timer
+            def inner_method():
+                return
+
+            inner_method()
+            return
+
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            outer_method()
+        lines = err.getvalue().strip().splitlines()
+        self.assertEqual(len(lines), 2)
+
+        # Inner finishes first (printed first)
+        self.assertRegex(
+            lines[0], rf"^ic\| inner_method took {TIMER_DURATION_RE_DECORATOR}$"
+        )
+        self.assertRegex(
+            lines[1], rf"^ic\| outer_method took {TIMER_DURATION_RE_DECORATOR}$"
+        )
+
+    def test_timer_decorator_return_value(self):
+        @ic.timer
+        def add(x: int, y: int) -> int:
+            return x + y
+
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            result: int = add(a, b)
+        self.assertEqual(result, a + b)
+        self.assertEqual(out.getvalue(), "")
+        pattern = rf"^ic\| add took {TIMER_DURATION_RE_DECORATOR}$"
+        self.assertRegex(err.getvalue().strip(), pattern)
+
+    def test_timer_decorator_with_arguments(self):
+        @ic.timer
+        def add(x, y=0):
+            return x + y
+
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            result = add(a, y=b)
+        self.assertEqual(result, a + b)
+        self.assertEqual(out.getvalue(), "")
+        pattern = rf"^ic\| add took {TIMER_DURATION_RE_DECORATOR}$"
+        self.assertRegex(err.getvalue().strip(), pattern)
+
+    def test_timer_decorator_when_disabled(self):
+        try:
+            ic.disable()
+            with capture_standard_streams() as (out, err):
+                noop_with_time_decorator()
+        finally:
+            ic.enable()
+        self.assertEqual(err.getvalue(), "")
+        self.assertEqual(out.getvalue(), "")
+
+    def test_timer_decorator_output_on_exception(self):
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            with self.assertRaises(ValueError):
+                raise_value_error()
+        pattern = (
+            rf"^ic\| {raise_value_error.__name__} took {TIMER_DURATION_RE_DECORATOR}$"
+        )
+        self.assertRegex(err.getvalue().strip(), pattern)
+
+    def test_timer_decorator_prefix_configuration(self):
+        prefix = "timer> "
+        with configure_icecream_output(prefix=prefix, outputFunction=stderr_print):
+            with capture_standard_streams() as (out, err):
+                noop_with_time_decorator()
+        pattern = rf"^timer> {noop_with_time_decorator.__name__} took {TIMER_DURATION_RE_DECORATOR}$"
+        self.assertRegex(err.getvalue().strip(), pattern)
+
+    def test_timer_decorator_output_function(self):
+        lst = []
+        with configure_icecream_output(outputFunction=lambda s: lst.append(s)):
+            noop_with_time_decorator()
+        self.assertEqual(len(lst), 1)
+        pattern = rf"^ic\| {noop_with_time_decorator.__name__} took {TIMER_DURATION_RE_DECORATOR}$"
+        self.assertRegex(lst[0], pattern)
+
+    def test_timer_decorator_preserves_function_metadata(self):
+        def original():
+            """My docstring."""
+            pass
+
+        wrapped = ic.timer(original)
+        self.assertEqual(wrapped.__name__, original.__name__)
+        self.assertEqual(wrapped.__doc__, original.__doc__)
+
+    def test_timer_context_manager_basic(self):
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            with ic.timer:
+                noop()
+
+        self.assertRegex(err.getvalue(), TIMER_DURATION_RE_CONTEXT_MANAGER)
+
+    def test_timer_context_manager_measures_elapsed_time(self):
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            with ic.timer:
+                time.sleep(0.1)
+                noop()
+        match = re.search(TIMER_DURATION_RE_CONTEXT_MANAGER, err.getvalue())
+
+        self.assertIsNotNone(match)
+
+        value, unit = float(match.group(1)), match.group(2)
+
+        multiplier = {"ns": 0.000001, "us": 0.001, "ms": 1, "s": 1000}
+        duration_ms = value * multiplier[unit]
+
+        self.assertGreater(duration_ms, 50)
+
+    def test_timer_context_manager_propagates_exception(self):
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            with self.assertRaises(ValueError):
+                with ic.timer:
+                    raise ValueError("Raised Error in timer block.")
+        self.assertRegex(err.getvalue(), TIMER_DURATION_RE_CONTEXT_MANAGER)
+
+    def test_timer_context_manager_exit_without_enter(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            ic.timer.__exit__(None, None, None)
+        self.assertIn("__enter__", str(ctx.exception))
+
+    def test_timer_context_manager_when_disabled(self):
+        try:
+            ic.disable()
+            with disable_coloring(), capture_standard_streams() as (out, err):
+                with ic.timer:
+                    noop()
+        finally:
+            ic.enable()
+        self.assertEqual(out.getvalue(), "")
+        self.assertEqual(err.getvalue(), "")
+
+    def test_timer_context_manager_resets_enter_time_after_exit(self):
+        timer = ic.timer
+        with disable_coloring(), capture_standard_streams() as (out, err):
+            with timer:
+                noop()
+
+        self.assertIsNone(timer._enter_time)
+        with self.assertRaises(RuntimeError) as ctx:
+            timer.__exit__(None, None, None)
+        self.assertIn("__enter__", str(ctx.exception))