-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathprofiler.py
More file actions
264 lines (225 loc) · 8.58 KB
/
Copy pathprofiler.py
File metadata and controls
264 lines (225 loc) · 8.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
"""Profile a target under :mod:`cProfile` + :mod:`tracemalloc`, emit a JSON report.
The report schema is intentionally small and stable so that two reports produced
on different machines/checkouts can be matched function-by-function by
:mod:`profile_diff.diff`.
Report schema (``schema_version = 1``)::
{
"schema_version": 1,
"label": "base", # from --label
"metadata": {
"python_version": "3.11.4",
"timestamp": "2026-07-18T12:00:00Z",
"target": "pytest [tests/]",
"top_n": 30
},
"total_time_s": 1.234, # total tottime across all functions
"functions": [
{"key": "pkg/mod.py:12(func)",
"ncalls": 10, "tottime": 0.5, "cumtime": 1.1},
... # top-N by cumtime, descending
],
"memory": {
"peak_bytes": 12345678,
"top_allocations": [
{"location": "pkg/mod.py:20", "size_bytes": 4096},
...
]
}
}
Function keys use the form ``file:line(func)``. File paths are normalised to be
relative to the current working directory when possible, so that the same source
file profiled from two different checkout directories yields the *same* key.
"""
from __future__ import annotations
import cProfile
import datetime as _dt
import importlib
import os
import pstats
import runpy
import sys
import tracemalloc
from typing import Callable, Dict, List, Optional, Tuple
from . import SCHEMA_VERSION
def normalize_path(filename: str) -> str:
"""Return a stable, checkout-independent form of *filename*.
Paths inside the current working directory are made relative to it (so
``/home/runner/base/pkg/mod.py`` and ``/home/runner/head/pkg/mod.py`` both
collapse to ``pkg/mod.py``). Synthetic names such as ``<string>`` and paths
outside the tree (e.g. the stdlib) are returned unchanged.
"""
if not filename or filename.startswith("<"):
return filename
try:
rel = os.path.relpath(filename, os.getcwd())
except ValueError:
# Different drive on Windows, etc.
return filename
# Only relativise when the file actually lives under cwd.
if not rel.startswith(os.pardir + os.sep) and rel != os.pardir:
return rel.replace(os.sep, "/")
return filename
def _func_key(filename: str, lineno: int, funcname: str) -> str:
return f"{normalize_path(filename)}:{lineno}({funcname})"
# Directory of this package; frames from here are profiler scaffolding, not the
# user's target, so they are excluded from reports to keep the signal on the
# code under test.
_PKG_DIR = os.path.dirname(os.path.abspath(__file__))
def _is_internal(filename: str) -> bool:
"""True for frames belonging to profile_diff's own harness code."""
if not filename or filename.startswith("<"):
return False
try:
return os.path.abspath(filename).startswith(_PKG_DIR + os.sep)
except (ValueError, OSError): # pragma: no cover - exotic paths
return False
# --------------------------------------------------------------------------- #
# Target runners
# --------------------------------------------------------------------------- #
def _run_pytest(pytest_args: List[str]) -> None:
"""Run ``python -m pytest`` in-process (so it is captured by cProfile)."""
argv = ["pytest", *pytest_args]
old_argv = sys.argv
sys.argv = argv
try:
# pytest's __main__ calls sys.exit(); swallow it so profiling completes.
runpy.run_module("pytest", run_name="__main__", alter_sys=True)
except SystemExit:
pass
finally:
sys.argv = old_argv
def _run_script(path: str, script_args: List[str]) -> None:
"""Execute a Python script by path as ``__main__``."""
argv = [path, *script_args]
old_argv = sys.argv
sys.argv = argv
try:
runpy.run_path(path, run_name="__main__")
finally:
sys.argv = old_argv
def _resolve_callable(ref: str) -> Callable[[], object]:
"""Resolve a ``pkg.mod:func`` (or ``pkg.mod:obj.method``) reference."""
module_name, sep, attr_path = ref.partition(":")
if not sep or not module_name or not attr_path:
raise ValueError(
f"--callable must look like 'pkg.mod:func', got {ref!r}"
)
obj = importlib.import_module(module_name)
for attr in attr_path.split("."):
obj = getattr(obj, attr)
if not callable(obj):
raise TypeError(f"{ref!r} is not callable")
return obj # type: ignore[return-value]
def _run_callable(ref: str) -> None:
_resolve_callable(ref)()
# --------------------------------------------------------------------------- #
# Core profiling
# --------------------------------------------------------------------------- #
def _collect_functions(profiler: cProfile.Profile, top_n: int) -> Tuple[List[Dict], float]:
"""Extract per-function records (top *top_n* by cumtime) and total time."""
stats = pstats.Stats(profiler)
records: List[Dict] = []
for (filename, lineno, funcname), (_cc, nc, tt, ct, _callers) in stats.stats.items(): # type: ignore[attr-defined]
if _is_internal(filename):
continue
records.append(
{
"key": _func_key(filename, lineno, funcname),
"ncalls": nc,
"tottime": tt,
"cumtime": ct,
}
)
# Deterministic ordering: cumtime desc, then key asc.
records.sort(key=lambda r: (-r["cumtime"], r["key"]))
total_time = float(getattr(stats, "total_tt", 0.0))
return records[:top_n], total_time
def _collect_memory(snapshot: tracemalloc.Snapshot, peak: int, top_n: int) -> Dict:
"""Build the memory section from a tracemalloc snapshot."""
stats = snapshot.statistics("lineno")
top: List[Dict] = []
for stat in stats:
frame = stat.traceback[0]
if _is_internal(frame.filename):
continue
if len(top) >= top_n:
break
top.append(
{
"location": f"{normalize_path(frame.filename)}:{frame.lineno}",
"size_bytes": int(stat.size),
}
)
return {"peak_bytes": int(peak), "top_allocations": top}
def profile(
run_fn: Callable[[], None],
*,
label: str,
target: str,
top_n: int = 30,
) -> Dict:
"""Profile ``run_fn`` and return a report dict following the schema above.
CPU stats come from :mod:`cProfile`; peak memory and top allocations come
from :mod:`tracemalloc`, both captured over the same execution.
"""
tracemalloc.start()
profiler = cProfile.Profile()
profiler.enable()
try:
run_fn()
finally:
profiler.disable()
_current, peak = tracemalloc.get_traced_memory()
snapshot = tracemalloc.take_snapshot()
tracemalloc.stop()
functions, total_time = _collect_functions(profiler, top_n)
memory = _collect_memory(snapshot, peak, top_n)
return {
"schema_version": SCHEMA_VERSION,
"label": label,
"metadata": {
"python_version": ".".join(str(v) for v in sys.version_info[:3]),
"timestamp": _dt.datetime.now(_dt.timezone.utc)
.replace(microsecond=0)
.isoformat()
.replace("+00:00", "Z"),
"target": target,
"top_n": top_n,
},
"total_time_s": total_time,
"functions": functions,
"memory": memory,
}
def build_run_fn(
*,
pytest_args: Optional[List[str]] = None,
script: Optional[str] = None,
script_args: Optional[List[str]] = None,
callable_ref: Optional[str] = None,
) -> Tuple[Callable[[], None], str]:
"""Return a ``(run_fn, target_description)`` pair for the selected target.
Exactly one of *pytest_args*, *script*, or *callable_ref* must be provided.
"""
selected = [
name
for name, value in (
("pytest", pytest_args),
("script", script),
("callable", callable_ref),
)
if value is not None
]
if len(selected) != 1:
raise ValueError(
"Exactly one target is required: --pytest, --script, or --callable"
)
if pytest_args is not None:
args = list(pytest_args)
desc = "pytest [" + " ".join(args) + "]" if args else "pytest"
return (lambda: _run_pytest(args)), desc
if script is not None:
args = list(script_args or [])
desc = "script " + script + (" " + " ".join(args) if args else "")
return (lambda: _run_script(script, args)), desc
assert callable_ref is not None
return (lambda: _run_callable(callable_ref)), "callable " + callable_ref