diff --git a/stopwatch/__init__.py b/stopwatch/__init__.py index e53007e..0ee894a 100644 --- a/stopwatch/__init__.py +++ b/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'] diff --git a/stopwatch/profile.py b/stopwatch/profile.py index da9753c..3f2c91c 100644 --- a/stopwatch/profile.py +++ b/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 diff --git a/stopwatch/statistics.py b/stopwatch/statistics.py index cb3a806..203e37f 100644 --- a/stopwatch/statistics.py +++ b/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) diff --git a/stopwatch/stopwatch.py b/stopwatch/stopwatch.py index 481301c..ae8c117 100644 --- a/stopwatch/stopwatch.py +++ b/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', ] ) diff --git a/stopwatch/utils.py b/stopwatch/utils.py index 0ff8d99..62a3a2f 100644 --- a/stopwatch/utils.py +++ b/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 '', - function=stack.function, - line_number=stack.lineno, + module=frame.f_globals.get('__name__', ''), + 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' diff --git a/tests/mocks/inspect.py b/tests/mocks/inspect.py deleted file mode 100644 index d28fcae..0000000 --- a/tests/mocks/inspect.py +++ /dev/null @@ -1,17 +0,0 @@ -from dataclasses import dataclass - - -@dataclass -class StackMock: - frame: str - function: str - lineno: int - - -class InspectMock: - @staticmethod - def stack() -> list[StackMock]: - return [ - StackMock(frame='frame', function='function', lineno=2) - for _ in range(3) - ] diff --git a/tests/test_profile.py b/tests/test_profile.py index e15c2bb..4277cd7 100644 --- a/tests/test_profile.py +++ b/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) diff --git a/tests/test_statistics.py b/tests/test_statistics.py index 40ab187..4d7a0db 100644 --- a/tests/test_statistics.py +++ b/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(), diff --git a/tests/test_stopwatch.py b/tests/test_stopwatch.py index 1d17dc1..6718076 100644 --- a/tests/test_stopwatch.py +++ b/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() diff --git a/tests/test_utils.py b/tests/test_utils.py index 922d130..ea31f4e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,10 +1,8 @@ +import inspect from unittest import TestCase -from unittest.mock import patch from stopwatch import utils -from .mocks.inspect import InspectMock - class UtilsTest(TestCase): def test_format_elapsed_time(self) -> None: @@ -16,9 +14,35 @@ def test_format_elapsed_time(self) -> None: self.assertEqual(utils.format_elapsed_time(0.0001, 0), '100μs') self.assertEqual(utils.format_elapsed_time(0.000001, 0), '1μs') - def test_inspect_caller(self) -> None: - with patch('inspect.stack', InspectMock.stack): - caller = utils.inspect_caller() - self.assertEqual(caller.module, '') - self.assertEqual(caller.function, 'function') - self.assertEqual(caller.line_number, 2) + def test_format_elapsed_time_picks_unit_by_magnitude(self) -> None: + self.assertEqual(utils.format_elapsed_time(-3600, 2), '-3600.00s') + self.assertEqual(utils.format_elapsed_time(-1, 2), '-1.00s') + self.assertEqual(utils.format_elapsed_time(-0.1, 2), '-100.00ms') + self.assertEqual(utils.format_elapsed_time(-0.000001, 2), '-1.00μs') + self.assertEqual(utils.format_elapsed_time(0, 2), '0.00μs') + + def test_inspect_caller_describes_two_frames_up(self) -> None: + def helper() -> utils.Caller: + return utils.inspect_caller() + + frame = inspect.currentframe() + assert frame is not None + expected_line = frame.f_lineno + 1 + caller = helper() + + self.assertEqual(caller.module, __name__) + self.assertEqual( + caller.function, 'test_inspect_caller_describes_two_frames_up' + ) + self.assertEqual(caller.line_number, expected_line) + + def test_inspect_caller_offset_skips_extra_frames(self) -> None: + def inner() -> utils.Caller: + return utils.inspect_caller(offset=1) + + def outer() -> utils.Caller: + return inner() + + self.assertEqual( + outer().function, 'test_inspect_caller_offset_skips_extra_frames' + )