Coverage for pyodmongo/models/db_decimal.py: 100%

1from typing import Any 

2from bson import Decimal128 

3from decimal import Decimal, Context, localcontext 

4from pydantic import GetCoreSchemaHandler 

5from pydantic_core import core_schema, CoreSchema 

6from pydantic_core.core_schema import ValidationInfo, str_schema, float_schema 

7 

8 

9class DbDecimal(Decimal): 

10 def __new__(cls, value: Any, scale: None | int = None): 

11 if isinstance(value, Decimal128): 

12 value = value.to_decimal() 

13 elif isinstance(value, (int, float, str, Decimal)): 

14 _DEC128_CTX = Context(prec=34) 

15 with localcontext(_DEC128_CTX) as ctx: 

16 value = ctx.create_decimal(value) 

17 if scale: 

18 quantizer = Decimal("0." + ("0" * scale)) 

19 value = value.quantize(quantizer) 

20 else: 

21 raise ValueError("Invalid decimal type") 

22 return super().__new__(cls, value) 

23 

24 def to_scale(self, scale: int): 

25 quantizer = Decimal("0." + ("0" * scale)) 

26 return self.quantize(quantizer) 

27 

28 @classmethod 

29 def validate(cls, v, _: ValidationInfo): 

30 return cls(v) 

31 

32 @classmethod 

33 def __get_pydantic_core_schema__( 

34 cls, source_type: Any, handler: GetCoreSchemaHandler 

35 ) -> CoreSchema: # type: ignore 

36 return core_schema.json_or_python_schema( 

37 python_schema=core_schema.with_info_plain_validator_function(cls.validate), 

38 json_schema=float_schema(), 

39 serialization=core_schema.plain_serializer_function_ser_schema( 

40 lambda v: float(v) if isinstance(v, Decimal) else v, when_used="json" 

41 ), 

42 )