codekingpro/portable-devtools

1from __future__ import annotations2 3import inspect4import logging5import types6import typing7 8from ._models import Request9 10 11class Trace:12    def __init__(13        self,14        name: str,15        logger: logging.Logger,16        request: Request | None = None,17        kwargs: dict[str, typing.Any] | None = None,18    ) -> None:19        self.name = name20        self.logger = logger21        self.trace_extension = (22            None if request is None else request.extensions.get("trace")23        )24        self.debug = self.logger.isEnabledFor(logging.DEBUG)25        self.kwargs = kwargs or {}26        self.return_value: typing.Any = None27        self.should_trace = self.debug or self.trace_extension is not None28        self.prefix = self.logger.name.split(".")[-1]29 30    def trace(self, name: str, info: dict[str, typing.Any]) -> None:31        if self.trace_extension is not None:32            prefix_and_name = f"{self.prefix}.{name}"33            ret = self.trace_extension(prefix_and_name, info)34            if inspect.iscoroutine(ret):  # pragma: no cover35                raise TypeError(36                    "If you are using a synchronous interface, "37                    "the callback of the `trace` extension should "38                    "be a normal function instead of an asynchronous function."39                )40 41        if self.debug:42            if not info or "return_value" in info and info["return_value"] is None:43                message = name44            else:45                args = " ".join([f"{key}={value!r}" for key, value in info.items()])46                message = f"{name} {args}"47            self.logger.debug(message)48 49    def __enter__(self) -> Trace:50        if self.should_trace:51            info = self.kwargs52            self.trace(f"{self.name}.started", info)53        return self54 55    def __exit__(56        self,57        exc_type: type[BaseException] | None = None,58        exc_value: BaseException | None = None,59        traceback: types.TracebackType | None = None,60    ) -> None:61        if self.should_trace:62            if exc_value is None:63                info = {"return_value": self.return_value}64                self.trace(f"{self.name}.complete", info)65            else:66                info = {"exception": exc_value}67                self.trace(f"{self.name}.failed", info)68 69    async def atrace(self, name: str, info: dict[str, typing.Any]) -> None:70        if self.trace_extension is not None:71            prefix_and_name = f"{self.prefix}.{name}"72            coro = self.trace_extension(prefix_and_name, info)73            if not inspect.iscoroutine(coro):  # pragma: no cover74                raise TypeError(75                    "If you're using an asynchronous interface, "76                    "the callback of the `trace` extension should "77                    "be an asynchronous function rather than a normal function."78                )79            await coro80 81        if self.debug:82            if not info or "return_value" in info and info["return_value"] is None:83                message = name84            else:85                args = " ".join([f"{key}={value!r}" for key, value in info.items()])86                message = f"{name} {args}"87            self.logger.debug(message)88 89    async def __aenter__(self) -> Trace:90        if self.should_trace:91            info = self.kwargs92            await self.atrace(f"{self.name}.started", info)93        return self94 95    async def __aexit__(96        self,97        exc_type: type[BaseException] | None = None,98        exc_value: BaseException | None = None,99        traceback: types.TracebackType | None = None,100    ) -> None:101        if self.should_trace:102            if exc_value is None:103                info = {"return_value": self.return_value}104                await self.atrace(f"{self.name}.complete", info)105            else:106                info = {"exception": exc_value}107                await self.atrace(f"{self.name}.failed", info)108