Fix profile measurements and make the instrument cheap

stopwatch/__init__.py

@@ -1,6 +1,10 @@
+from colorama import just_fix_windows_console
+
 from .profile import profile
 from .statistics import Statistics
 from .stopwatch import Stopwatch
 from .utils import format_elapsed_time
 
+just_fix_windows_console()
+
 __all__ = ['Statistics', 'Stopwatch', 'format_elapsed_time', 'profile']

stopwatch/profile.py

@@ -1,8 +1,8 @@
 import atexit
 import functools
-import math
+import inspect
 from collections.abc import Callable
-from typing import Any, TypeVar
+from typing import Any, TypeVar, cast
 
 from colorama import Fore, Style
 
@@ -47,7 +47,7 @@ def _make_report(caller: Caller, name: str, statistics: Statistics) -> str:
  f'min={format_elapsed_time(statistics.minimum)}',
  f'median={format_elapsed_time(statistics.median)}',
  f'max={format_elapsed_time(statistics.maximum)}',
- f'dev={format_elapsed_time(math.sqrt(statistics.variance))}',
+ f'dev={format_elapsed_time(statistics.stdev)}',
  ]
  )
 
@@ -71,37 +71,83 @@ def _print_report(caller: Caller, name: str, statistics: Statistics) -> None:
  print(_make_report(caller, name, statistics))
 
 
-def profile(**kwargs: Any) -> Callable[[Callable[..., RT]], Callable[..., RT]]:
+def profile(
+ *,
+ name: str | None = None,
+ report_every: int | None = 1,
+) -> Callable[[Callable[..., RT]], Callable[..., RT]]:
  """
- Decorator for profiling the function.
+ Decorator for profiling the function. Must be called: `@profile()`.
+
+ Works on regular and `async def` functions. Generator functions are
+ not supported -- see the note below.
 
  Parameters
  ----------
  name : Optional[`str`]
  The name for the statistics. Default is the name of function.
  report_every : Optional[`int`]
- The number of times to report the statistics. Default is 1.
+ Report once every this many calls, or `None` to only report when
+ the process exits. Default is 1.
+
+ Raises
+ ------
+ `ValueError`
+ If `report_every` is smaller than 1.
+
+ Notes
+ -----
+ A generator function is measured only for building the generator, not
+ for consuming it. Every recorded value is retained until the process
+ exits, since the median needs all of them.
  """
+ if report_every is not None and report_every < 1:
+ raise ValueError(
+ f'report_every must be >= 1 or None, got {report_every}'
+ )
+
  caller = inspect_caller()
 
  def decorator(func: Callable[..., RT]) -> Callable[..., RT]:
- name: str = kwargs.get('name', func.__name__)
- report_every: int = kwargs.get('report_every', 1)
- should_report = report_every is not None
-
+ stat_name = func.__name__ if name is None else name
  statistics = Statistics()
- atexit.register(_print_report, caller, name, statistics)
 
- @functools.wraps(func)
- def wrapper(*args: object, **kwargs: object) -> RT:
- with Stopwatch() as stopwatch:
- result = func(*args, **kwargs)
+ def is_report_due() -> bool:
+ return (
+ report_every is not None
+ and len(statistics) % report_every == 0
+ )
+
+ def record(elapsed: float) -> None:
+ statistics.add(elapsed)
+ if is_report_due():
+ _print_report(caller, stat_name, statistics)
 
- statistics.add(stopwatch.elapsed)
- if should_report and (len(statistics) % report_every) == 0:
- _print_report(caller, name, statistics)
+def report_at_exit() -> None:
+ if not is_report_due():
+ _print_report(caller, stat_name, statistics)
 
- return result
+ atexit.register(report_at_exit)
+
+ if inspect.iscoroutinefunction(func):
+
+ @functools.wraps(func)
+ async def async_wrapper(*args: object, **kwargs: object) -> Any:
+ stopwatch = Stopwatch()
+ try:
+ return await func(*args, **kwargs)
+ finally:
+ record(stopwatch.stop().elapsed)
+
+ return cast('Callable[..., RT]', async_wrapper)
+
+ @functools.wraps(func)
+ def wrapper(*args: object, **kwargs: object) -> RT:
+ stopwatch = Stopwatch()
+ try:
+ return func(*args, **kwargs)
+ finally:
+ record(stopwatch.stop().elapsed)
 
  return wrapper
 

stopwatch/statistics.py

@@ -5,7 +5,7 @@ class Statistics:
  _values: list[float]
 
  def __init__(self, values: list[float] | None = None) -> None:
- self._values = values or []
+ self._values = list(values) if values else []
 
  def add(self, value: float) -> None:
  """
@@ -41,13 +41,18 @@ def minimum(self) -> float:
  @property
  def total(self) -> float:
  """`float`: Return the total value in seconds."""
- return sum(self._values)
+ return float(sum(self._values))
 
  @property
  def variance(self) -> float:
  """`float`: Return the variance in seconds."""
  return statistics.pvariance(self._values)
 
+ @property
+ def stdev(self) -> float:
+ """`float`: Return the population standard deviation in seconds."""
+ return statistics.pstdev(self._values)
+
  def __len__(self) -> int:
  return len(self._values)
 

stopwatch/stopwatch.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import math
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Any
@@ -176,7 +175,7 @@ def report(self) -> str:
  f'min={statistics.minimum:.{self.precision}f}s',
  f'median={statistics.median:.{self.precision}f}s',
  f'max={statistics.maximum:.{self.precision}f}s',
- f'dev={math.sqrt(statistics.variance):.{self.precision}f}s',
+ f'dev={statistics.stdev:.{self.precision}f}s',
  ]
  )
 

stopwatch/utils.py

@@ -1,4 +1,4 @@
-import inspect
+import sys
 from typing import NamedTuple
 
 
@@ -9,12 +9,24 @@ class Caller(NamedTuple):
 
 
 def inspect_caller(offset: int = 0) -> Caller:
- stack = inspect.stack()[2 + offset]
- module = inspect.getmodule(stack.frame)
+ """
+ Describe the frame two levels above this call.
+
+ Parameters
+ ----------
+ offset : `int`
+ Extra frames to skip, for callers that add a level of their own.
+
+ Returns
+ -------
+ `Caller`
+ The module name, function name and line number of that frame.
+ """
+ frame = sys._getframe(2 + offset)
  return Caller(
- module=module.__name__ if module else '<unknown>',
- function=stack.function,
- line_number=stack.lineno,
+ module=frame.f_globals.get('__name__', '<unknown>'),
+ function=frame.f_code.co_name,
+ line_number=frame.f_lineno,
  )
 
 
@@ -35,8 +47,9 @@ def format_elapsed_time(elapsed: float, precision: int = 2) -> str:
  The formatted elapsed time.
  """
  ms = elapsed * 1e3
- if ms >= 1e3:
+ magnitude = abs(ms)
+ if magnitude >= 1e3:
  return f'{ms / 1e3:.{precision}f}s'
- if ms >= 1:
+ if magnitude >= 1:
  return f'{ms:.{precision}f}ms'
  return f'{ms * 1e3:.{precision}f}μs'

tests/test_profile.py

@@ -1,3 +1,6 @@
+import asyncio
+import sys
+from collections.abc import Callable
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
@@ -109,3 +112,83 @@ def test_profile_with_report_every_none(
  decorated_func()
  self.assertEqual(func.call_count, 4)
  print_mock.assert_not_called()
+
+ def test_profile_rejects_invalid_report_every(self) -> None:
+ with self.assertRaises(ValueError):
+ profile(report_every=0)
+ with self.assertRaises(ValueError):
+ profile(report_every=-1)
+
+ @patch('builtins.print')
+ def test_profile_records_calls_that_raise(
+ self, print_mock: MagicMock
+ ) -> None:
+ registered: list[object] = []
+ with patch('atexit.register', registered.append):
+ with patch('time.perf_counter', self.time_mock.perf_counter):
+
+ @profile(name='boom')
+ def failing() -> None:
+ self.time_mock.increment(0.5)
+ raise RuntimeError('nope')
+
+ for _ in range(2):
+ with self.assertRaises(RuntimeError):
+ failing()
+ self.assertEqual(print_mock.call_count, 2)
+ report = print_mock.call_args[0][0]
+ self.assertIn('hits=2', report)
+ self.assertIn('mean=500.00ms', report)
+
+ def test_profile_measures_async_functions(self) -> None:
+ totals: list[float] = []
+ profile_module = sys.modules['stopwatch.profile']
+ with (
+ patch('atexit.register', lambda *a, **k: None),
+ patch.object(
+ profile_module,
+ '_print_report',
+ lambda caller, name, stats: totals.append(stats.total),
+ ),
+ ):
+
+ @profile(name='coro')
+ async def work() -> str:
+ await asyncio.sleep(0.05)
+ return 'done'
+
+ self.assertEqual(asyncio.run(work()), 'done')
+
+ self.assertEqual(len(totals), 1)
+ self.assertGreaterEqual(totals[0], 0.05)
+
+ @patch('builtins.print')
+ def test_profile_does_not_report_twice_at_exit(
+ self, print_mock: MagicMock
+ ) -> None:
+ registered: list[Callable[[], None]] = []
+ with patch('atexit.register', registered.append):
+ with patch('time.perf_counter', self.time_mock.perf_counter):
+ func = MagicMock(__name__='dup')
+ decorated_func = profile()(func)
+ self.time_mock.increment(0.1)
+ decorated_func()
+ self.assertEqual(print_mock.call_count, 1)
+ registered[0]()
+ self.assertEqual(print_mock.call_count, 1)
+
+ @patch('builtins.print')
+ def test_profile_flushes_unreported_calls_at_exit(
+ self, print_mock: MagicMock
+ ) -> None:
+ registered: list[Callable[[], None]] = []
+ with patch('atexit.register', registered.append):
+ with patch('time.perf_counter', self.time_mock.perf_counter):
+ func = MagicMock(__name__='flush')
+ decorated_func = profile(report_every=2)(func)
+ for _ in range(3):
+ self.time_mock.increment(0.1)
+ decorated_func()
+ self.assertEqual(print_mock.call_count, 1)
+ registered[0]()
+ self.assertEqual(print_mock.call_count, 2)

tests/test_statistics.py

@@ -46,6 +46,20 @@ def test_variance(self) -> None:
  self.stats.variance, statistics.pvariance(self.values)
  )
 
+ def test_stdev(self) -> None:
+ self.assertEqual(self.stats.stdev, statistics.pstdev(self.values))
+
+ def test_does_not_alias_the_given_list(self) -> None:
+ values = [1.0, 2.0]
+ stats = Statistics(values)
+ stats.add(3.0)
+ self.assertEqual(values, [1.0, 2.0])
+ values.append(99.0)
+ self.assertEqual(len(stats), 3)
+
+ def test_total_is_a_float_when_empty(self) -> None:
+ self.assertIsInstance(Statistics().total, float)
+
  def test_to_dict(self) -> None:
  self.assertEqual(
  self.stats.to_dict(),

tests/test_stopwatch.py

@@ -146,6 +146,16 @@ def test_stopwatch_without_autostart_still_starts_as_context(self) -> None:
  self.assertTrue(sw.running)
  self.assertEqual(sw.elapsed, 1.0)
 
+ def test_stopwatch_start_while_running_is_a_noop(self) -> None:
+ with patch('time.perf_counter', self.time_mock.perf_counter):
+ sw = Stopwatch()
+ self.time_mock.increment(1)
+ sw.start()
+ self.time_mock.increment(1)
+ sw.stop()
+ self.assertEqual(len(sw.laps), 1)
+ self.assertEqual(sw.elapsed, 2.0)
+
  def test_stopwatch_lap_closes_on_exception(self) -> None:
  with patch('time.perf_counter', self.time_mock.perf_counter):
  sw = Stopwatch()