Coverage for polar/custom_field/data.py: 79%

1import datetime 

2from collections.abc import Sequence 

3from typing import TYPE_CHECKING, Any 

4 

5from pydantic import BaseModel, Field, ValidationError, create_model 

6from sqlalchemy import event 

7from sqlalchemy.dialects.postgresql import JSONB 

8from sqlalchemy.orm import Mapped, Mapper, ORMDescriptor, mapped_column 

9 

10from polar.exceptions import PolarRequestValidationError 

11 

12if TYPE_CHECKING: 12 ↛ 13line 12 didn't jump to line 13 because the condition on line 12 was never true

13 from polar.models import CustomField, Organization 

14 

15 from .attachment import AttachedCustomFieldMixin 

16 

17 

18class CustomFieldDataMixin: 

19 custom_field_data: Mapped[ 

20 dict[str, str | int | bool | datetime.datetime | None] 

21 ] = mapped_column(JSONB, nullable=False, default=dict) 

22 

23 # Make the type checker happy, but we should make sure actual models 

24 # declare an organization attribute by themselves. 

25 if TYPE_CHECKING: 25 ↛ anywhereline 25 didn't jump anywhere: it always raised an exception.

26 organization: ORMDescriptor["Organization"] 

27 

28 

29custom_field_data_models: set[type[CustomFieldDataMixin]] = set() 

30 

31 

32# Event listener to track models inheriting from CustomFieldDataMixin 

33@event.listens_for(Mapper, "mapper_configured") 

34def track_attached_custom_field_mixin(_mapper: Mapper[Any], class_: type) -> None: 

35 if issubclass(class_, CustomFieldDataMixin): 

36 custom_field_data_models.add(class_) 

37 

38 

39class CustomFieldDataInputMixin(BaseModel): 

40 custom_field_data: dict[str, str | int | bool | datetime.datetime | None] = Field( 

41 default_factory=dict, 

42 description="Key-value object storing custom field values.", 

43 ) 

44 

45 

46class CustomFieldDataOutputMixin(BaseModel): 

47 custom_field_data: dict[str, str | int | bool | datetime.datetime | None] = Field( 

48 default_factory=dict, 

49 description="Key-value object storing custom field values.", 

50 ) 

51 

52 

53def build_custom_field_data_schema( 

54 custom_fields: Sequence[tuple["CustomField", bool]], 

55) -> type[BaseModel]: 

56 fields_definitions: Any = { 

57 custom_field.slug: custom_field.get_field_definition(required) 

58 for custom_field, required in custom_fields 

59 } 

60 return create_model("CustomFieldDataInput", **fields_definitions) 

61 

62 

63def validate_custom_field_data( 

64 attached_custom_fields: Sequence["AttachedCustomFieldMixin"], 

65 data: dict[str, Any], 

66 *, 

67 error_loc_prefix: Sequence[str] = ("body", "custom_field_data"), 

68 validate_required: bool = True, 

69) -> dict[str, Any]: 

70 schema = build_custom_field_data_schema( 

71 [ 

72 (f.custom_field, validate_required and f.required) 

73 for f in attached_custom_fields 

74 ] 

75 ) 

76 try: 

77 return schema.model_validate(data).model_dump(mode="json") 

78 except ValidationError as e: 

79 raise PolarRequestValidationError( 

80 [{**err, "loc": (*error_loc_prefix, *err["loc"])} for err in e.errors()] # pyright: ignore 

81 )