43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
from __future__ import annotations
|
|
from contextlib import ContextDecorator, contextmanager
|
|
from typing import Any, Generator
|
|
|
|
from .node_profiler import NodeProfiler, get_current_profiler, set_current_profiler
|
|
|
|
|
|
@contextmanager
|
|
def profiler_context() -> Generator[NodeProfiler, None, None]:
|
|
"""Context manager that instantiates and activates a NodeProfiler for tree execution."""
|
|
old = get_current_profiler()
|
|
profiler = NodeProfiler()
|
|
set_current_profiler(profiler)
|
|
try:
|
|
yield profiler
|
|
finally:
|
|
set_current_profiler(old)
|
|
|
|
|
|
class NodeProfileExecution(ContextDecorator):
|
|
"""Context manager and decorator that profiles a node's execution if a profiler is active,
|
|
or acts as a zero-cost pass-through if no profiler is active."""
|
|
|
|
def __init__(self, node: Any):
|
|
self.node = node
|
|
self._ctx = None
|
|
|
|
def __enter__(self):
|
|
profiler = get_current_profiler()
|
|
if profiler is not None:
|
|
self._ctx = profiler.profile_node(self.node)
|
|
self._ctx.__enter__()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if self._ctx is not None:
|
|
return self._ctx.__exit__(exc_type, exc_val, exc_tb)
|
|
return False
|
|
|
|
|
|
def profile_node_execution(node: Any) -> NodeProfileExecution:
|
|
"""Profiles the execution of node if a profiler context is active."""
|
|
return NodeProfileExecution(node)
|