orjson.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """JSON Formatter using [orjson](https://github.com/ijl/orjson)"""
  2. ### IMPORTS
  3. ### ============================================================================
  4. ## Future
  5. from __future__ import annotations
  6. ## Standard Library
  7. from typing import Any
  8. from collections.abc import Callable
  9. ## Installed
  10. ## Application
  11. from . import core
  12. from . import defaults as d
  13. from .utils import package_is_available
  14. # We import orjson after checking it is available
  15. package_is_available("orjson", throw_error=True)
  16. import orjson # pylint: disable=wrong-import-position,wrong-import-order
  17. ### FUNCTIONS
  18. ### ============================================================================
  19. def orjson_default(obj: Any) -> Any:
  20. """orjson default encoder function for non-standard types"""
  21. if d.use_exception_default(obj):
  22. return d.exception_default(obj)
  23. if d.use_traceback_default(obj):
  24. return d.traceback_default(obj)
  25. if d.use_bytes_default(obj):
  26. return d.bytes_default(obj)
  27. if d.use_enum_default(obj):
  28. return d.enum_default(obj)
  29. if d.use_type_default(obj):
  30. return d.type_default(obj)
  31. return d.unknown_default(obj)
  32. ### CLASSES
  33. ### ============================================================================
  34. class OrjsonFormatter(core.BaseJsonFormatter):
  35. """JSON formatter using [orjson](https://github.com/ijl/orjson) for encoding."""
  36. def __init__(
  37. self,
  38. *args,
  39. json_default: Callable | None = orjson_default,
  40. json_indent: bool = False,
  41. **kwargs,
  42. ) -> None:
  43. """
  44. Args:
  45. args: see [BaseJsonFormatter][pythonjsonlogger.core.BaseJsonFormatter]
  46. json_default: a function for encoding non-standard objects
  47. json_indent: indent output with 2 spaces.
  48. kwargs: see [BaseJsonFormatter][pythonjsonlogger.core.BaseJsonFormatter]
  49. """
  50. super().__init__(*args, **kwargs)
  51. self.json_default = json_default
  52. self.json_indent = json_indent
  53. return
  54. def jsonify_log_record(self, log_data: core.LogData) -> str:
  55. """Returns a json string of the log data."""
  56. opt = orjson.OPT_NON_STR_KEYS
  57. if self.json_indent:
  58. opt |= orjson.OPT_INDENT_2
  59. return orjson.dumps(log_data, default=self.json_default, option=opt).decode("utf8")