import inspect import sys from collections import Counter from contextlib import AbstractContextManager, ExitStack from functools import partial from io import StringIO from itertools import count from types import CodeType, ModuleType from typing import (Any, Optional, Union, Callable, Generator, Dict, FrozenSet, Tuple, ClassVar) import pytest from line_profiler import _line_profiler, LineProfiler USE_SYS_MONITORING = isinstance(getattr(sys, 'monitoring', None), ModuleType) # -------------------------------------------------------------------- # # Helper classes # # -------------------------------------------------------------------- # class SysMonHelper: """ Helper object which helps with simplifying attribute access on :py:mod:`sys.monitoring`. """ tool_id: int no_tool_id_callables: ClassVar[FrozenSet[str]] = frozenset( {'restart_events'}) def __init__(self, tool_id: Optional[int] = None) -> None: if tool_id is None: tool_id = sys.monitoring.PROFILER_ID self.tool_id = tool_id def __getattr__(self, attr: str): """ Returns: * If ``attr`` refers to a :py:mod:`sys.monitoring` callable: a :py:func:`functools.partial` object with :py:attr:`~.tool_id` pre-applied, unless it is in the set of explicitly excluded methods (:py:attr:`~.no_tool_id_callables`). * If ``attr`` is all uppercase: the corresponding integer or special (e.g. :py:data:`~sys.monitoring.MISSING` and :py:data:`~sys.monitoring.DISABLE`) constants from either the module or the :py:data:`sys.monitoring.events` namespace. * Otherwise, it falls through to the module. """ mon = sys.monitoring if attr.isupper(): try: return getattr(mon.events, attr) except AttributeError: pass result = getattr(mon, attr) if callable(result) and attr not in self.no_tool_id_callables: return partial(result, self.tool_id) return result def get_current_callback( self, event_id: Optional[int] = None) -> Union[Callable, None]: """ Arguments: event_id (int | None) Optional integer ID to retrieve the callback from; defaults to :py:data:`sys.monitoring.events.LINE`. Returns: The current callback (if any) associated with with :py:data:`~.tool_id` and ``event_id``. """ register = self.register_callback if event_id is None: event_id = MON.LINE result = register(event_id, None) if result is not None: register(event_id, result) return result class restore_events(AbstractContextManager): """ Restore the global or local :py:mod:`sys.monitoring` events. """ code: Union[CodeType, None] mon: SysMonHelper events: int def __init__(self, *, code: Optional[CodeType] = None, tool_id: Optional[int] = None) -> None: self.code = code self.mon = SysMonHelper(tool_id) self.events = sys.monitoring.events.NO_EVENTS def __enter__(self): if self.code is None: self.events = self.mon.get_events() else: self.events = self.mon.get_local_events(self.code) return self def __exit__(self, *_, **__) -> None: if self.mon.get_tool() is None: pass # No-op elif self.code is None: self.mon.set_events(self.events) else: self.mon.set_local_events(self.code, self.events) self.events = self.mon.NO_EVENTS class LineCallback: """ Simple :py:mod:`sys.monitoring` callback for handling LINE events. Attributes: nhits (dict[tuple[str, int, str], Counter[int]]) Mapping from :py:attr:`line_profiler._line_profiler.LineStats.timings` keys to a :py:class:`collections.Counter` mapping line numbers to reported hit counts. predicate (Callable[[code, int], bool]) Callable taking the code object and line number, and returning whether the line event should be reported. disable (bool) Settable boolean determining whether to return :py:data:`sys.monitoring.DISABLE` on a reported line event. """ nhits: Dict[Tuple[str, int, str], 'Counter[int]'] predicate: Callable[[CodeType, int], bool] disable: bool def __init__( self, predicate: Callable[[CodeType, int], bool], *, register: bool = True, disable: bool = False ) -> None: """ Arguments: predicate, disable See attributes. register If true, register the instance with :py:func:`sys.monitoring.register_callback`. """ self.nhits = {} self.predicate = predicate self.disable = disable if register: self.register() def register(self) -> Any: """ Returns: Old value of the :py.mod:`sys.monitoring` callback. Side effects: Instance registered as the new :py:mod:`sys.monitoring` callback. """ return MON.register_callback(MON.LINE, self) def handle_line_event(self, code: CodeType, lineno: int) -> bool: """ Returns: Whether a line event has been recorded. Side effects: Entry created/incremented in :py:attr:`~.nhits` if :py:attr:`~.predicate` evaluates to true. """ result = self.predicate(code, lineno) if result: self.nhits.setdefault( _line_profiler.label(code), Counter())[lineno] += 1 return result def __call__(self, code: CodeType, lineno: int) -> Any: """ Returns: :py:data:`sys.monitoring.DISABLE` if :py:attr:`~.predicate` evaluates to true AND if :py:attr:`~.disable` is true; :py:const:`None` otherwise. Side effects: Entry created/incremented in :py:attr:`~.nhits` if :py:attr:`~.predicate` evaluates to true. """ if self.handle_line_event(code, lineno) and self.disable: return MON.DISABLE if USE_SYS_MONITORING: MON = SysMonHelper() # -------------------------------------------------------------------- # # Helper functions # # -------------------------------------------------------------------- # def enable_line_events(code: Optional[CodeType] = None) -> None: if code is None: MON.set_events(MON.get_events() | MON.LINE) else: MON.set_local_events(code, MON.get_local_events(code) | MON.LINE) def disable_line_events(code: Optional[CodeType] = None) -> None: if code is None: events = MON.get_events() if events & MON.LINE: MON.set_events(events ^ MON.LINE) else: events = MON.get_local_events(code) if events & MON.LINE: MON.set_local_events(code, events ^ MON.LINE) # -------------------------------------------------------------------- # # Tests # # -------------------------------------------------------------------- # @pytest.fixture(autouse=True) def sys_mon_cleanup( monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: """ If :py:mod:`sys.monitoring` is available: * Make sure that we are using the default behavior by overriding :py:data:`line_profiler._line_profiler.USE_LEGACY_TRACE`; * Remember all the relevant global callbacks before running the test; * If ``sys.monitoring.PROFILER_ID`` isn't already in use, mark it as being used by `'line_profiler_tests'`; and * Finally, restore these after the test: - The callbacks - The tool name (if we have set it earlier) - The globally-active events (or the lack thereof) for ``sys.monitoring.PROFILER_ID`` otherwise, automatically :py:func:`pytest.skip` the test. """ def restore(message): for name, callback in callbacks.items(): prev_callback = MON.register_callback(event_ids[name], callback) if prev_callback is callback: callback_repr = '(UNCHANGED)' else: callback_repr = '-> ' + repr(callback) print('{} (`sys.monitoring.events.{}`): {!r} {}'.format( message, name, prev_callback, callback_repr)) if not USE_SYS_MONITORING: pytest.skip('No `sys.monitoring`') # Remember the callbacks event_ids = {name: getattr(MON, name) for name in ('LINE', 'PY_RETURN', 'PY_YIELD')} callbacks = {name: MON.register_callback(event_id, None) for name, event_id in event_ids.items()} # Restore the callbacks since we "popped" them restore('Pre-test: putting the callbacks back') # Set the tool name if it isn't set already set_tool_id = not MON.get_tool() if set_tool_id: MON.use_tool_id('line_profiler_tests') try: monkeypatch.setattr(_line_profiler, 'USE_LEGACY_TRACE', False) with restore_events(): yield finally: # Restore the callbacks restore('Post-test: restoring the callbacks') # Unset the temporary tool name if set_tool_id: MON.free_tool_id() def test_standalone_callback_usage() -> None: """ Check that :py:mod:`sys.monitoring` callbacks behave as expected when `LineProfiler`s are not in use. """ _test_callback_helper(6, 7, 8, 9) @pytest.mark.parametrize('wrap_trace', [True, False]) def test_wrapping_trace(wrap_trace: bool) -> None: """ Check that existing :py:mod:`sys.monitoring` callbacks behave as expected depending on `LineProfiler.wrap_trace`. """ prof = LineProfiler(wrap_trace=wrap_trace) try: nhits_expected = _test_callback_helper( 6, 7, 8, 9, prof=prof, callback_called=wrap_trace) finally: with StringIO() as sio: prof.print_stats(sio) output = sio.getvalue() print(output) line = next(line for line in output.splitlines() if line.endswith('# Loop body')) nhits = int(line.split()[1]) assert nhits == nhits_expected def _test_callback_helper( nloop_no_trace: int, nloop_trace_global: int, nloop_trace_local: int, nloop_disabled: int, prof: Optional[LineProfiler] = None, callback_called: bool = True) -> int: cumulative_nhits = 0 def func(n: int) -> int: x = 0 for n in range(1, n + 1): x += n # Loop body return x def get_loop_hits() -> int: nonlocal cumulative_nhits cumulative_nhits = ( callback.nhits[_line_profiler.label(code)][lineno_loop]) return cumulative_nhits lines, first_lineno = inspect.getsourcelines(func) lineno_loop = first_lineno + next( offset for offset, line in enumerate(lines) if line.rstrip().endswith('# Loop body')) names = {func.__name__, func.__qualname__} code = func.__code__ if prof is not None: orig_func, func = func, prof(func) code = orig_func.__code__ callback = LineCallback(lambda code, _: code.co_name in names) # When line events are suppressed, nothing should happen with restore_events(): disable_line_events() n = nloop_no_trace assert MON.get_current_callback() is callback assert func(n) == n * (n + 1) // 2 assert MON.get_current_callback() is callback assert not callback.nhits # When line events are activated, the callback should see line # events with restore_events(): # Global events enable_line_events() n = nloop_trace_global assert MON.get_current_callback() is callback assert func(n) == n * (n + 1) // 2 assert MON.get_current_callback() is callback print(callback.nhits) if callback_called: expected = cumulative_nhits + n assert get_loop_hits() == expected else: assert not callback.nhits with ExitStack() as stack: stack.enter_context(restore_events()) stack.enter_context(restore_events(code=code)) # Disable global line events, and enable local line events disable_line_events() enable_line_events(code) n = nloop_trace_local assert MON.get_current_callback() is callback assert func(n) == n * (n + 1) // 2 assert MON.get_current_callback() is callback print(callback.nhits) if callback_called: expected = cumulative_nhits + n assert get_loop_hits() == expected else: assert not callback.nhits # Line events can be disabled on the specific line by returning # `sys.monitoring.DISABLE` from the callback for enable_global in True, False: callback.disable = True with ExitStack() as stack: stack.enter_context(restore_events()) stack.enter_context(restore_events(code=code)) # Set the global and local events # (doesn't matter if events are enabled globally or not) if enable_global: enable_line_events() else: disable_line_events() enable_line_events(code) # We still get 1 more hit because that's the call we # return `sys.monitoring.DISABLE` from n = nloop_disabled assert MON.get_current_callback() is callback assert func(n) == n * (n + 1) // 2 assert MON.get_current_callback() is callback print(callback.nhits) if callback_called: expected = cumulative_nhits + 1 assert get_loop_hits() == expected else: assert not callback.nhits MON.restart_events() # Return the total number of loops run # (Note: `nloop_disabled` is used twice) return (nloop_no_trace + nloop_trace_global + nloop_trace_local + 2 * nloop_disabled) @pytest.mark.parametrize('standalone', [True, False]) def test_callback_switching(standalone: bool) -> None: """ Check that hot-swapping of :py:mod:`sys.monitoring` callbacks behaves as expected, no matter if `LineProfiler`s are in use. """ if standalone: prof: Union[LineProfiler, None] = None else: prof = LineProfiler(wrap_trace=True) try: nhits_expected = _test_callback_switching_helper(17, prof) finally: if prof is not None: with StringIO() as sio: prof.print_stats(sio) output = sio.getvalue() print(output) if prof is None: return line = next(line for line in output.splitlines() if line.endswith('# Loop body')) nhits = int(line.split()[1]) assert nhits == nhits_expected def _test_callback_switching_helper( nloop: int, prof: Optional[LineProfiler] = None) -> int: cumulative_nhits = 0, 0 def func(n: int) -> int: x = 0 for n in range(1, n + 1): x += n # Loop body return x def get_loop_hits() -> Tuple[int, int]: nonlocal cumulative_nhits cumulative_nhits = tuple( # type: ignore[assignment] callback.nhits.get( _line_profiler.label(code), Counter())[lineno_loop] for callback in (callback_1, callback_2)) return cumulative_nhits def predicate(code: CodeType, lineno: int) -> bool: return code.co_name in names and lineno == lineno_loop class SwitchingCallback(LineCallback): """ Callback which switches to the next one after having been triggered. """ next: Union['SwitchingCallback', None] def __init__(self, *args, next: Optional['SwitchingCallback'] = None, **kwargs) -> None: super().__init__(*args, **kwargs) self.next = next def __call__(self, code: CodeType, lineno: int) -> Any: if not self.handle_line_event(code, lineno): return if self.next is not None: self.next.register() if self.disable: return MON.DISABLE lines, first_lineno = inspect.getsourcelines(func) lineno_loop = first_lineno + next( offset for offset, line in enumerate(lines) if line.rstrip().endswith('# Loop body')) names = {func.__name__, func.__qualname__} code = func.__code__ if prof is not None: orig_func, func = func, prof(func) code = orig_func.__code__ # The two callbacks hand off to one another in a loop callback_1 = SwitchingCallback(predicate) callback_2 = SwitchingCallback(predicate, register=False, next=callback_1) callback_1.next = callback_2 with restore_events(): enable_line_events() assert MON.get_current_callback() is callback_1 assert func(nloop) == nloop * (nloop + 1) // 2 assert MON.get_current_callback() in (callback_1, callback_2) print(callback_1.nhits, callback_2.nhits) nhits_one = nloop // 2 nhits_other = nloop - nhits_one if nhits_one == nhits_other: assert get_loop_hits() == (nhits_one, nhits_other) else: # Odd number assert get_loop_hits() in ((nhits_one, nhits_other), (nhits_other, nhits_one)) return nloop @pytest.mark.parametrize('add_events', [True, False]) @pytest.mark.parametrize('code_local_events', [True, False]) @pytest.mark.parametrize('start_with_events', [True, False]) @pytest.mark.parametrize('standalone', [True, False]) def test_callback_update_events( standalone: bool, start_with_events: bool, code_local_events: bool, add_events: bool) -> None: """ Check that a :py:mod:`sys.monitoring` callback which updates the event set (global and code-object-local) after a certain number of hits behaves as expected, no matter if `LineProfiler`s are in use. """ nloop = 10 nloop_update = 5 cumulative_nhits = 0 def func(n: int) -> int: x = 0 for n in range(1, n + 1): x += n # Loop body return x def get_loop_hits() -> int: nonlocal cumulative_nhits cumulative_nhits = ( callback.nhits[_line_profiler.label(code)][lineno_loop]) return cumulative_nhits class EventUpdatingCallback(LineCallback): """ Callback which, after a certain number of hits: - Disables :py:attr:`sys.monitoring.LINE` events, and - Enables :py:attr:`sys.monitoring.CALL` events (if :py:attr:`~.call` is true) """ def __init__(self, *args, code: Optional[CodeType] = None, call: bool = False, **kwargs) -> None: super().__init__(*args, **kwargs) self.count = 0 self.code = code self.call = call def __call__(self, code: CodeType, lineno: int) -> None: if not self.handle_line_event(code, lineno): return self.count += 1 if self.count >= nloop_update: disable_line_events(self.code) if not self.call: return if self.code is None: MON.set_events(MON.get_events() | MON.CALL) else: MON.set_local_events( self.code, MON.get_local_events(self.code) | MON.CALL) lines, first_lineno = inspect.getsourcelines(func) lineno_loop = first_lineno + next( offset for offset, line in enumerate(lines) if line.rstrip().endswith('# Loop body')) names = {func.__name__, func.__qualname__} code = func.__code__ if standalone: prof: Union[LineProfiler, None] = None else: prof = LineProfiler(wrap_trace=True) if start_with_events: global_events = MON.CALL else: global_events = MON.NO_EVENTS if prof is not None: orig_func, func = func, prof(func) code = orig_func.__code__ callback = EventUpdatingCallback( lambda code, lineno: (code.co_name in names and lineno == lineno_loop), code=code if code_local_events else None, call=add_events) local_events = local_events_after = MON.NO_EVENTS global_events_after = global_events if code_local_events: local_events |= MON.LINE if add_events: local_events_after |= MON.CALL else: global_events |= MON.LINE if add_events: global_events_after |= MON.CALL MON.set_events(global_events) MON.set_local_events(code, local_events) try: # Check that data is only gathered for the first `nloop_update` # hits assert MON.get_current_callback() is callback assert func(nloop) == nloop * (nloop + 1) // 2 assert MON.get_current_callback() is callback print(callback.nhits) assert get_loop_hits() == nloop_update # Check that the callback has disabled LINE events (and enabled # CALL events where appropriate) assert MON.get_events() == global_events_after assert MON.get_local_events(code) == local_events_after finally: if prof is not None: with StringIO() as sio: prof.print_stats(sio) output = sio.getvalue() print(output) if prof is None: return line = next(line for line in output.splitlines() if line.endswith('# Loop body')) nhits = int(line.split()[1]) assert nhits == nloop @pytest.mark.parametrize('standalone', [True, False]) def test_callback_toggle_local_events(standalone: bool) -> None: """ Check that a :py:mod:`sys.monitoring` callback which disables local LINE events and later re-enables them with :py:func:`sys.monitoring.restart_events` behaves as expected, no matter if `LineProfiler`s are in use. """ if standalone: prof: Union[LineProfiler, None] = None else: prof = LineProfiler(wrap_trace=True) try: nhits_expected = _test_callback_toggle_local_events_helper( 17, 18, 19, prof) finally: if prof is not None: with StringIO() as sio: prof.print_stats(sio) output = sio.getvalue() print(output) if prof is None: return line = next(line for line in output.splitlines() if line.endswith('# Loop body')) nhits = int(line.split()[1]) assert nhits == nhits_expected def _test_callback_toggle_local_events_helper( nloop_before_disabling: int, nloop_when_disabled: int, nloop_after_reenabling: int, prof: Optional[LineProfiler] = None) -> int: cumulative_nhits = 0 def func(*nloops) -> int: x = 0 counter = count(1) for n in nloops: for _ in range(n): x += next(counter) # Loop body pass # Switching location return x def get_loop_hits() -> int: nonlocal cumulative_nhits cumulative_nhits = ( callback.nhits[_line_profiler.label(code)][lineno_loop]) return cumulative_nhits class LocalDisablingCallback(LineCallback): """ Callback which disables LINE events locally after a certain number of hits """ def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.switch_count = 0 def __call__(self, code: CodeType, lineno: int) -> Any: if not self.handle_line_event(code, lineno): return # When we hit the "loop" line after having hit the "switch" # line, disable line events on the "loop" line if lineno == lineno_loop and self.switch_count % 2: # Remove the recorded hit self.nhits[_line_profiler.label(code)][lineno] -= 1 return MON.DISABLE # When we hit the "switch" line the second line, restart # events to undo disabling of the "loop" line if lineno == lineno_switch: if self.switch_count % 2: MON.restart_events() self.switch_count += 1 return lines, first_lineno = inspect.getsourcelines(func) lineno_loop = first_lineno + next( offset for offset, line in enumerate(lines) if line.rstrip().endswith('# Loop body')) lineno_switch = first_lineno + next( offset for offset, line in enumerate(lines) if line.rstrip().endswith('# Switching location')) linenos = {lineno_loop, lineno_switch} names = {func.__name__, func.__qualname__} code = func.__code__ if prof is not None: orig_func, func = func, prof(func) code = orig_func.__code__ callback = LocalDisablingCallback( lambda code, lineno: (code.co_name in names and lineno in linenos)) MON.set_events(MON.get_events() | MON.LINE) n = nloop_before_disabling + nloop_when_disabled + nloop_after_reenabling assert MON.get_current_callback() is callback assert func(nloop_before_disabling, nloop_when_disabled, nloop_after_reenabling) == n * (n + 1) // 2 assert MON.get_current_callback() is callback print(callback.nhits) assert get_loop_hits() == nloop_before_disabling + nloop_after_reenabling return n