You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

154 lines
5.3 KiB

  1. import dataclasses
  2. from collections import defaultdict
  3. from enum import Enum
  4. from pathlib import PurePath
  5. from types import GeneratorType
  6. from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
  7. from pydantic import BaseModel
  8. from pydantic.json import ENCODERS_BY_TYPE
  9. SetIntStr = Set[Union[int, str]]
  10. DictIntStrAny = Dict[Union[int, str], Any]
  11. def generate_encoders_by_class_tuples(
  12. type_encoder_map: Dict[Any, Callable[[Any], Any]]
  13. ) -> Dict[Callable[[Any], Any], Tuple[Any, ...]]:
  14. encoders_by_class_tuples: Dict[Callable[[Any], Any], Tuple[Any, ...]] = defaultdict(
  15. tuple
  16. )
  17. for type_, encoder in type_encoder_map.items():
  18. encoders_by_class_tuples[encoder] += (type_,)
  19. return encoders_by_class_tuples
  20. encoders_by_class_tuples = generate_encoders_by_class_tuples(ENCODERS_BY_TYPE)
  21. def jsonable_encoder(
  22. obj: Any,
  23. include: Optional[Union[SetIntStr, DictIntStrAny]] = None,
  24. exclude: Optional[Union[SetIntStr, DictIntStrAny]] = None,
  25. by_alias: bool = True,
  26. exclude_unset: bool = False,
  27. exclude_defaults: bool = False,
  28. exclude_none: bool = False,
  29. custom_encoder: Dict[Any, Callable[[Any], Any]] = {},
  30. sqlalchemy_safe: bool = True,
  31. ) -> Any:
  32. if include is not None and not isinstance(include, (set, dict)):
  33. include = set(include)
  34. if exclude is not None and not isinstance(exclude, (set, dict)):
  35. exclude = set(exclude)
  36. if isinstance(obj, BaseModel):
  37. encoder = getattr(obj.__config__, "json_encoders", {})
  38. if custom_encoder:
  39. encoder.update(custom_encoder)
  40. obj_dict = obj.dict(
  41. include=include, # type: ignore # in Pydantic
  42. exclude=exclude, # type: ignore # in Pydantic
  43. by_alias=by_alias,
  44. exclude_unset=exclude_unset,
  45. exclude_none=exclude_none,
  46. exclude_defaults=exclude_defaults,
  47. )
  48. if "__root__" in obj_dict:
  49. obj_dict = obj_dict["__root__"]
  50. return jsonable_encoder(
  51. obj_dict,
  52. exclude_none=exclude_none,
  53. exclude_defaults=exclude_defaults,
  54. custom_encoder=encoder,
  55. sqlalchemy_safe=sqlalchemy_safe,
  56. )
  57. if dataclasses.is_dataclass(obj):
  58. return dataclasses.asdict(obj)
  59. if isinstance(obj, Enum):
  60. return obj.value
  61. if isinstance(obj, PurePath):
  62. return str(obj)
  63. if isinstance(obj, (str, int, float, type(None))):
  64. return obj
  65. if isinstance(obj, dict):
  66. encoded_dict = {}
  67. for key, value in obj.items():
  68. if (
  69. (
  70. not sqlalchemy_safe
  71. or (not isinstance(key, str))
  72. or (not key.startswith("_sa"))
  73. )
  74. and (value is not None or not exclude_none)
  75. and ((include and key in include) or not exclude or key not in exclude)
  76. ):
  77. encoded_key = jsonable_encoder(
  78. key,
  79. by_alias=by_alias,
  80. exclude_unset=exclude_unset,
  81. exclude_none=exclude_none,
  82. custom_encoder=custom_encoder,
  83. sqlalchemy_safe=sqlalchemy_safe,
  84. )
  85. encoded_value = jsonable_encoder(
  86. value,
  87. by_alias=by_alias,
  88. exclude_unset=exclude_unset,
  89. exclude_none=exclude_none,
  90. custom_encoder=custom_encoder,
  91. sqlalchemy_safe=sqlalchemy_safe,
  92. )
  93. encoded_dict[encoded_key] = encoded_value
  94. return encoded_dict
  95. if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
  96. encoded_list = []
  97. for item in obj:
  98. encoded_list.append(
  99. jsonable_encoder(
  100. item,
  101. include=include,
  102. exclude=exclude,
  103. by_alias=by_alias,
  104. exclude_unset=exclude_unset,
  105. exclude_defaults=exclude_defaults,
  106. exclude_none=exclude_none,
  107. custom_encoder=custom_encoder,
  108. sqlalchemy_safe=sqlalchemy_safe,
  109. )
  110. )
  111. return encoded_list
  112. if custom_encoder:
  113. if type(obj) in custom_encoder:
  114. return custom_encoder[type(obj)](obj)
  115. else:
  116. for encoder_type, encoder in custom_encoder.items():
  117. if isinstance(obj, encoder_type):
  118. return encoder(obj)
  119. if type(obj) in ENCODERS_BY_TYPE:
  120. return ENCODERS_BY_TYPE[type(obj)](obj)
  121. for encoder, classes_tuple in encoders_by_class_tuples.items():
  122. if isinstance(obj, classes_tuple):
  123. return encoder(obj)
  124. errors: List[Exception] = []
  125. try:
  126. data = dict(obj)
  127. except Exception as e:
  128. errors.append(e)
  129. try:
  130. data = vars(obj)
  131. except Exception as e:
  132. errors.append(e)
  133. raise ValueError(errors)
  134. return jsonable_encoder(
  135. data,
  136. by_alias=by_alias,
  137. exclude_unset=exclude_unset,
  138. exclude_defaults=exclude_defaults,
  139. exclude_none=exclude_none,
  140. custom_encoder=custom_encoder,
  141. sqlalchemy_safe=sqlalchemy_safe,
  142. )