def _instrument_aiohttp()

in aidial_sdk/header_propagator.py [0:0]

• 21 lines of code
• 3 McCabe index (conditional complexity)

    def _instrument_aiohttp(self):
        try:
            import aiohttp
        except ImportError:
            return

        async def _on_request_start(
            session: aiohttp.ClientSession,
            trace_config_ctx: types.SimpleNamespace,
            params: aiohttp.TraceRequestStartParams,
        ):
            self._modify_headers(str(params.url), params.headers)

        def instrumented_init(wrapped, instance, args, kwargs):
            trace_config = aiohttp.TraceConfig()
            trace_config.on_request_start.append(_on_request_start)

            trace_configs = list(kwargs.get("trace_configs") or [])
            trace_configs.append(trace_config)

            kwargs["trace_configs"] = trace_configs
            return wrapped(*args, **kwargs)

        wrapt.wrap_function_wrapper(
            aiohttp.ClientSession, "__init__", instrumented_init
        )