Doc/tools/extensions/profiling_trace.py

"""
Sphinx extension to generate profiler trace data during docs build.

This extension executes a demo Python program with sys.settrace() to capture
the execution trace and injects it into the profiling visualization JS file.
"""

import json
import re
import sys
from io import StringIO
from pathlib import Path

from sphinx.errors import ExtensionError

DEMO_SOURCE = """\
def add(a, b):
    return a + b

def multiply(x, y):
    result = 0
    for i in range(y):
        result = add(result, x)
    return result

def calculate(a, b):
    sum_val = add(a, b)
    product = multiply(a, b)
    return sum_val + product

def main():
    result = calculate(3, 4)
    print(f"Result: {result}")

main()
"""

PLACEHOLDER = "/* PROFILING_TRACE_DATA */null"


def generate_trace(source: str) -> list[dict]:
    """
    Execute the source code with tracing enabled and capture execution events.
    """
    trace_events = []
    timestamp = [0]
    timestamp_step = 50
    tracing_active = [False]
    pending_line = [None]

    def tracer(frame, event, arg):
        if frame.f_code.co_filename != '<demo>':
            return tracer

        func_name = frame.f_code.co_name
        lineno = frame.f_lineno

        if event == 'line' and not tracing_active[0]:
            pending_line[0] = {'type': 'line', 'line': lineno}
            return tracer

        # Start tracing only once main() is called
        if event == 'call' and func_name == 'main':
            tracing_active[0] = True
            # Emit the buffered line event (the main() call line) at ts=0
            if pending_line[0]:
                pending_line[0]['ts'] = 0
                trace_events.append(pending_line[0])
                pending_line[0] = None
                timestamp[0] = timestamp_step

        # Skip events until we've entered main()
        if not tracing_active[0]:
            return tracer

        if event == 'call':
            trace_events.append({
                'type': 'call',
                'func': func_name,
                'line': lineno,
                'ts': timestamp[0],
            })
        elif event == 'line':
            trace_events.append({
                'type': 'line',
                'line': lineno,
                'ts': timestamp[0],
            })
        elif event == 'return':
            try:
                value = arg if arg is None else repr(arg)
            except Exception:
                value = '<unprintable>'
            trace_events.append({
                'type': 'return',
                'func': func_name,
                'ts': timestamp[0],
                'value': value,
            })

            if func_name == 'main':
                tracing_active[0] = False

        timestamp[0] += timestamp_step
        return tracer

    # Suppress print output during tracing
    old_stdout = sys.stdout
    sys.stdout = StringIO()

    old_trace = sys.gettrace()
    sys.settrace(tracer)
    try:
        code = compile(source, '<demo>', 'exec')
        exec(code, {'__name__': '__main__'})
    finally:
        sys.settrace(old_trace)
        sys.stdout = old_stdout

    return trace_events


def inject_trace(app, exception):
    if exception:
        return

    js_path = (
        Path(app.outdir) / '_static' / 'profiling-sampling-visualization.js'
    )

    if not js_path.exists():
        return

    trace = generate_trace(DEMO_SOURCE)

    demo_data = {'source': DEMO_SOURCE.rstrip(), 'trace': trace, 'samples': []}

    demo_json = json.dumps(demo_data, indent=2)
    content = js_path.read_text(encoding='utf-8')

    pattern = r"(const DEMO_SIMPLE\s*=\s*/\* PROFILING_TRACE_DATA \*/)[^;]+;"

    if re.search(pattern, content):
        content = re.sub(
            pattern, lambda m: f"{m.group(1)} {demo_json};", content
        )
        js_path.write_text(content, encoding='utf-8')
        print(
            f"profiling_trace: Injected {len(trace)} trace events into {js_path.name}"
        )
    else:
        raise ExtensionError(
            f"profiling_trace: Placeholder pattern not found in {js_path.name}"
        )


def add_assets(app, pagename, templatename, context, doctree):
    if pagename == 'library/profiling.sampling':
        app.add_js_file('profiling-sampling-visualization.js')
        app.add_css_file('profiling-sampling-visualization.css')


def setup(app):
    app.connect('build-finished', inject_trace)
    app.connect('html-page-context', add_assets)

    return {
        'version': '1.0',
        'parallel_read_safe': True,
        'parallel_write_safe': True,
    }