1  from __future__ import annotations
   2  
   3  import inspect
   4  import logging
   5  import types
   6  import typing
   7  
   8  from ._models import Request
   9  
  10  
  11  class 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 = name
  20          self.logger = logger
  21          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 = None
  27          self.should_trace = self.debug or self.trace_extension is not None
  28          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 cover
  35                  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 = name
  44              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.kwargs
  52              self.trace(f"{self.name}.started", info)
  53          return self
  54  
  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 cover
  74                  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 coro
  80  
  81          if self.debug:
  82              if not info or "return_value" in info and info["return_value"] is None:
  83                  message = name
  84              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.kwargs
  92              await self.atrace(f"{self.name}.started", info)
  93          return self
  94  
  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