"""
AutoProfile Script Demo
=======================
The following demo is end-to-end bash code that writes a demo script and
profiles it with autoprofile.
.. code:: bash
# Write demo python script to disk
python -c "if 1:
import textwrap
text = textwrap.dedent(
'''
def plus(a, b):
return a + b
def fib(n):
a, b = 0, 1
while a < n:
a, b = b, plus(a, b)
def main():
import math
import time
start = time.time()
print('start calculating')
while time.time() - start < 1:
fib(10)
math.factorial(1000)
print('done calculating')
main()
'''
).strip()
with open('demo.py', 'w') as file:
file.write(text)
"
echo "---"
echo "## Profile With AutoProfile"
python -m kernprof -p demo.py -l demo.py
python -m line_profiler -rmt demo.py.lprof
"""
from __future__ import annotations
import importlib.util
import sys
import types
from collections.abc import MutableMapping
from typing import Any, cast, Dict, Mapping
from typing import ContextManager
from .ast_tree_profiler import AstTreeProfiler
from .run_module import AstTreeModuleProfiler
from .line_profiler_utils import add_imported_function_or_module
from .util_static import modpath_to_modname
PROFILER_LOCALS_NAME = 'prof'
def _extend_line_profiler_for_profiling_imports(prof: Any) -> None:
"""Allow profiler to handle functions/methods, classes & modules with a single call.
Add a method to LineProfiler that can identify whether the object is a
function/method, class or module and handle it's profiling accordingly.
Mainly used for profiling objects that are imported.
(Workaround to keep changes needed by autoprofile separate from base LineProfiler)
Args:
prof (LineProfiler):
instance of LineProfiler.
"""
prof.add_imported_function_or_module = types.MethodType(
add_imported_function_or_module, prof
)
[docs]
def run(
script_file: str,
ns: MutableMapping[str, Any],
prof_mod: list[str],
profile_imports: bool = False,
as_module: bool = False,
) -> None:
"""Automatically profile a script and run it.
Profile functions, classes & modules specified in prof_mod without needing to add
@profile decorators.
Args:
script_file (str):
path to script being profiled.
ns (dict):
"locals" from kernprof scope.
prof_mod (List[str]):
list of imports to profile in script.
passing the path to script will profile the whole script.
the objects can be specified using its dotted path or full path (if applicable).
profile_imports (bool):
if True, when auto-profiling whole script, profile all imports aswell.
as_module (bool):
Whether we're running script_file as a module
"""
class restore_dict:
def __init__(self, d: MutableMapping[str, Any]):
self.d = d
self.copy: Mapping[str, Any] | None = None
def __enter__(self):
assert self.copy is None
self.copy = dict(self.d)
def __exit__(self, *_, **__):
self.d.clear()
if self.copy is not None:
self.d.update(self.copy)
self.copy = None
Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler]
if as_module:
Profiler = AstTreeModuleProfiler
module_name = modpath_to_modname(script_file)
if not module_name:
raise ModuleNotFoundError(
f'script_file = {script_file!r}: cannot find corresponding module'
)
module_obj = types.ModuleType(module_name)
# Set the `__spec__` correctly
module_obj.__spec__ = importlib.util.find_spec(module_name)
else:
Profiler = AstTreeProfiler
module_obj = types.ModuleType('__main__')
namespace: MutableMapping[str, Any] = vars(module_obj)
namespace.update(ns)
profiler = Profiler(script_file, prof_mod, profile_imports)
tree_profiled = profiler.profile()
_extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME])
code_obj = compile(tree_profiled, script_file, 'exec')
with restore_dict(sys.modules):
# Always set the module object to `sys.modules['__main__']` and
# then restore it via the context manager, so that the executed
# code is run as `__main__`
sys.modules['__main__'] = module_obj
exec(code_obj, cast(Dict[str, Any], namespace), namespace) # type: ignore[redundant-cast]