Coverage for src/metador_core/schema/encoder.py: 100%

1"""Support for dynamically registered JSON encoders in pydantic. 

2 

3Preparation: 

4 

5For your top level BaseModel, set `DynEncoderModelMeta` as metaclass, e.g. 

6 

7``` 

8class MyBaseModel(BaseModel, metaclass=DynEncoderModelMeta): 

9 ... 

10``` 

11 

12(If you already use a custom metaclass for your base model, 

13add the `DynJsonEncoder` metaclass mixin) 

14 

15Usage: 

16 

17Decorate any class with `@json_encoder(ENCODER_FUNCTION)`. 

18 

19To add an encoder for some existing class you cannot decorate, 

20use `add_json_encoder(ClassName, func)`. 

21 

22Note that `json_encoders` declared as intended by Pydantic in the `Config` section 

23will always be prioritized over the dynamic encoder. This means, that the dynamic 

24encoders are only triggered for classes that are not models themselves 

25(because pydantic handles them already). 

26 

27Also note that to prevent bugs, you cannot override encoders for a class that already 

28has a registered dynamic encoder. Use the normal pydantic mechanisms for 

29cases where this is really needed. 

30 

31Ideally, design your classes in a way that there is a 1-to-1 relationship between 

32class and desired JSON encoder (e.g. you can have different subclasses with 

33different encoders). 

34""" 

35 

36from typing import Any, Callable, Dict, Type 

37 

38from pydantic.main import ModelMetaclass 

39 

40_reg_json_encoders: Dict[Type, Callable[[Any], str]] = {} 

41"""Global registry of declared JSON encoders.""" 

42 

43 

44def json_encoder(func): 

45 """Decorate a class to register a new JSON encoder for it.""" 

46 

47 def reg_encoder(cls): 

48 if issubclass(cls.__class__, ModelMetaclass): 

49 raise TypeError("This decorator does not work for pydantic models!") 

50 if hasattr(cls, "__dataclass_fields__"): 

51 raise TypeError("This decorator does not work for dataclasses!") 

52 

53 if cls in _reg_json_encoders: 

54 raise ValueError(f"A JSON encoder function for {cls} already exists!") 

55 

56 _reg_json_encoders[cls] = func 

57 return cls 

58 

59 return reg_encoder 

60 

61 

62def add_json_encoder(cls, func): 

63 """Register a JSON encoder function for a class.""" 

64 return json_encoder(func)(cls) 

65 

66 

67# ---- 

68 

69 

70def _dynamize_encoder(encoder_func): 

71 """Wrap the JSON encoder pydantic generates from the Config to support the dynamic registry.""" 

72 

73 def wrapped_encoder(obj): 

74 try: 

75 # try the default lookup 

76 return encoder_func(obj) 

77 except TypeError as e: 

78 if enc := _reg_json_encoders.get(type(obj)): 

79 return enc(obj) # try dynamic lookup 

80 raise e 

81 

82 return wrapped_encoder 

83 

84 

85class DynJsonEncoderMetaMixin(type): 

86 """Metaclass mixin to first look in dynamic encoder registry. 

87 

88 Combine this with (a subclass of) `ModelMetaClass` and use it for your custom base model. 

89 """ 

90 

91 def __init__(self, name, bases, dct): 

92 super().__init__(name, bases, dct) 

93 self.__json_encoder__ = staticmethod(_dynamize_encoder(self.__json_encoder__)) 

94 

95 

96class DynEncoderModelMetaclass(DynJsonEncoderMetaMixin, ModelMetaclass): 

97 """Set this metaclass for your custom base model to enable dynamic encoders."""