_tracing.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """
  2. Tracing utils
  3. """
  4. from __future__ import annotations
  5. from collections.abc import Sequence
  6. from typing import Any
  7. from typing import Callable
  8. _Writer = Callable[[str], object]
  9. _Processor = Callable[[tuple[str, ...], tuple[Any, ...]], object]
  10. class TagTracer:
  11. def __init__(self) -> None:
  12. self._tags2proc: dict[tuple[str, ...], _Processor] = {}
  13. self._writer: _Writer | None = None
  14. self.indent = 0
  15. def get(self, name: str) -> TagTracerSub:
  16. return TagTracerSub(self, (name,))
  17. def _format_message(self, tags: Sequence[str], args: Sequence[object]) -> str:
  18. if isinstance(args[-1], dict):
  19. extra = args[-1]
  20. args = args[:-1]
  21. else:
  22. extra = {}
  23. content = " ".join(map(str, args))
  24. indent = " " * self.indent
  25. lines = ["{}{} [{}]\n".format(indent, content, ":".join(tags))]
  26. for name, value in extra.items():
  27. lines.append(f"{indent} {name}: {value}\n")
  28. return "".join(lines)
  29. def _processmessage(self, tags: tuple[str, ...], args: tuple[object, ...]) -> None:
  30. if self._writer is not None and args:
  31. self._writer(self._format_message(tags, args))
  32. try:
  33. processor = self._tags2proc[tags]
  34. except KeyError:
  35. pass
  36. else:
  37. processor(tags, args)
  38. def setwriter(self, writer: _Writer | None) -> None:
  39. self._writer = writer
  40. def setprocessor(self, tags: str | tuple[str, ...], processor: _Processor) -> None:
  41. if isinstance(tags, str):
  42. tags = tuple(tags.split(":"))
  43. else:
  44. assert isinstance(tags, tuple)
  45. self._tags2proc[tags] = processor
  46. class TagTracerSub:
  47. def __init__(self, root: TagTracer, tags: tuple[str, ...]) -> None:
  48. self.root = root
  49. self.tags = tags
  50. def __call__(self, *args: object) -> None:
  51. self.root._processmessage(self.tags, args)
  52. def get(self, name: str) -> TagTracerSub:
  53. return self.__class__(self.root, self.tags + (name,))