snek/src/snek/system/profiler.py

import cProfile
import pstats
import sys 
from aiohttp import web
profiler = None
import io 


@web.middleware
async def profile_middleware(request, handler):
    global profiler 
    if not profiler:
        profiler = cProfile.Profile()
    profiler.enable()
    response = await handler(request)
    profiler.disable()
    stats = pstats.Stats(profiler, stream=sys.stdout)
    stats.sort_stats('cumulative')
    stats.print_stats() 
    return response

async def profiler_handler(request):
    output = io.StringIO()
    stats = pstats.Stats(profiler, stream=output)
    sort_by = request.query.get("sort", "tot. percall")
    stats.sort_stats(sort_by)
    stats.print_stats()
    return web.Response(text=output.getvalue())

class Profiler:

    def __init__(self):
        global profiler 
        if profiler is None:
            profiler = cProfile.Profile()
        self.profiler = profiler

    async def __aenter__(self):
        self.profiler.enable()    
    
    async def __aexit__(self, *args, **kwargs):
        self.profiler.disable()